Merge branch 'main' into feat-add-orpo-support

Author: eghuzefa

examples/deepscaler/math_eval_nb.py

@@ -20,6 +20,7 @@ classifiers = [
 dependencies = [
   "datasets",
   "flax>=0.11.1",
+  "fsspec",  # gcsfs dependency
   "gcsfs",
   "grain",
   "huggingface_hub",
@@ -28,9 +29,11 @@ dependencies = [
   "kagglehub",
   "numba",
   "omegaconf", # CLI config
+  "pylatexenc",  # Eval result parsing
   "python-dotenv",  # Huggingface API key
   "qwix",
   "sentencepiece",
+  "sympy",  # Eval result parsing
   "tensorboardX",
   "tensorflow_datasets",
   "tqdm",

examples/deepscaler/train_deepscaler.ipynb

@@ -85,7 +85,10 @@
     type=int,
     default=1869,
     required=False,
-    help="Number of batches for training.",
+    help=(
+        "Number of batches for training. Defaults to total number of samples //"
+        " global batch size."
+    ),
 )
 parser.add_argument(
     "--num-test-batches",
@@ -94,6 +97,27 @@
     required=False,
     help="Number of test batches for evaluation.",
 )
+parser.add_argument(
+    "--global-batch-size",
+    type=int,
+    default=4,
+    required=False,
+    help="Number of global batches for learning.",
+)
+parser.add_argument(
+    "--train-micro-batch-size",
+    type=int,
+    default=2,
+    required=False,
+    help="Number of micro batches for training.",
+)
+parser.add_argument(
+    "--train-mini-batch-size",
+    type=int,
+    default=4,
+    required=False,
+    help="Number of mini batches for training.",
+)
 parser.add_argument(
     "--rollout-engine",
     type=str,
@@ -163,7 +187,7 @@
 # ====== GRPO ======
 # === Generation during GRPO training ===
 MAX_PROMPT_LENGTH = 256
-TOTAL_GENERATION_STEPS = 1024  # YY 768
+TOTAL_GENERATION_STEPS = 768
 # Important to keep a high-ish temperature for varied, diverse responses during
 # training.
 TEMPERATURE = 0.9
@@ -186,17 +210,14 @@
 EPSILON = 0.2
 
 # ====== Training ======
-# 2 is the max we can do on v5e-8 with llama3 8B model.
-# 4 is the max we can do on v5e-8 with llama3 1B model.
-TRAIN_MICRO_BATCH_SIZE = 4
 # To speed up for quick workflow validation, we can change NUM_BATCHES to e.g. 2
-NUM_BATCHES = args.num_batches
+NUM_BATCHES = min(args.num_batches, 7473 // args.global_batch_size)
 # Keep `NUM_TEST_BATCHES` low so that evaluation runs quickly. It can be
 # increased to a max. of 330 (if batch size is 4).
 # To speed up for quick workflow validation, we can change it to e.g. 1
 NUM_TEST_BATCHES = args.num_test_batches
 
-EVAL_EVERY_N_STEPS = 10  # this doesn't matter if `TRAIN_FRACTION = 1.0`.
+EVAL_EVERY_N_STEPS = 1000  # this doesn't matter if `TRAIN_FRACTION = 1.0`.
 NUM_EPOCHS = 1  # can potentially train for more epochs
 
 # Number of training steps.
@@ -344,7 +365,7 @@ def get_dataset(path: str) -> grain.MapDataset:
   return loaded_dataset
 
 
-dataset = get_dataset(TRAIN_DATA_PATH).batch(TRAIN_MICRO_BATCH_SIZE)[
+dataset = get_dataset(TRAIN_DATA_PATH).batch(args.global_batch_size)[
     :NUM_BATCHES
 ]
 
@@ -357,7 +378,7 @@ def get_dataset(path: str) -> grain.MapDataset:
 
   val_dataset = dataset[int(len(dataset) * TRAIN_FRACTION) :].repeat(NUM_EPOCHS)
 
-test_dataset = get_dataset(TEST_DATA_PATH).batch(TRAIN_MICRO_BATCH_SIZE)[
+test_dataset = get_dataset(TEST_DATA_PATH).batch(args.global_batch_size)[
     :NUM_TEST_BATCHES
 ]
 
@@ -627,7 +648,7 @@ def generate(
 
   out_data = sampler(
       input_strings=input_batch,
-      max_generation_steps=768,
+      max_generation_steps=TOTAL_GENERATION_STEPS,
       temperature=temperature,
       top_k=top_k,
       top_p=top_p,
@@ -782,8 +803,8 @@ def evaluate(
         actor_optimizer=optimizer,
         eval_every_n_steps=EVAL_EVERY_N_STEPS,
         max_steps=MAX_STEPS,
-        mini_batch_size=TRAIN_MICRO_BATCH_SIZE,
-        train_micro_batch_size=TRAIN_MICRO_BATCH_SIZE,
+        mini_batch_size=args.train_mini_batch_size,
+        train_micro_batch_size=args.train_micro_batch_size,
         # metrics logging
         metrics_logging_options=metrics_logging_options,
         # checkpoint saving
@@ -802,7 +823,6 @@ def evaluate(
         rollout_vllm_tpu_backend_type="jax",
         rollout_vllm_server_mode=args.rollout_server_mode,
     ),
-
 )
 
 grpo_config = grpo_learner.GRPOConfig(

pyproject.toml

@@ -162,7 +162,7 @@ def test_sglang_jax_sampler(self):
     self.assertTrue(
         np.allclose(
             tunix_state["embedder"]["input_embedding"].value,
-            sglangjax_state["transformer"]["embed_tokens"]["embedding"].value,
+            sglangjax_state["model"]["embed_tokens"]["embedding"].value,
         )
     )
 

scripts/grpo_demo_llama3_qwen2.py

@@ -108,7 +108,6 @@ def _sglang_jax_config(self, config: SglangJaxConfig):
     args["model_path"] = config.model_version
     args["precompile_bs_paddings"] = [1, 64]
     args["precompile_token_paddings"] = [8192]
-    args["disable_jax_precompile"] = True
     args["page_size"] = 64
     args["context_length"] = config.context_length
     args["tp_size"] = self._find_tp_size(config.mesh)

tests/generate/sglang_jax_sampler_test.py

@@ -13,46 +13,46 @@ def _to_sglang_jax_mappings() -> Dict[str, MappingEntry]:
   return {
       'lm_head.w': ('lm_head.embedding', (None, 'model')),
       'embedder.input_embedding': (
-          'transformer.embed_tokens.embedding',
+          'model.embed_tokens.embedding',
           ('model', None),
       ),
       'layers.*.input_layernorm.w': (
-          'transformer.layers.*.input_layernorm.scale',
+          'model.layers.*.input_layernorm.scale',
           (None,),
       ),
       'layers.*.mlp.down_proj.kernel': (
-          'transformer.layers.*.mlp.down_proj.weight',
+          'model.layers.*.mlp.down_proj.weight',
           ('model', None),
       ),
       'layers.*.mlp.gate_proj.kernel': (
-          'transformer.layers.*.mlp.gate_proj.weight',
+          'model.layers.*.mlp.gate_proj.weight',
           (None, 'model'),
       ),
       'layers.*.mlp.up_proj.kernel': (
-          'transformer.layers.*.mlp.up_proj.weight',
+          'model.layers.*.mlp.up_proj.weight',
           (None, 'model'),
       ),
       'layers.*.post_attention_layernorm.w': (
-          'transformer.layers.*.post_attention_layernorm.scale',
+          'model.layers.*.post_attention_layernorm.scale',
           (None,),
       ),
       'layers.*.attn.k_proj.w': (
-          'transformer.layers.*.self_attn.k_proj.weight',
+          'model.layers.*.self_attn.k_proj.weight',
           (None, 'model', None),
       ),
       'layers.*.attn.o_proj.w': (
-          'transformer.layers.*.self_attn.o_proj.weight',
+          'model.layers.*.self_attn.o_proj.weight',
           ('model', None, None),
       ),
       'layers.*.attn.q_proj.w': (
-          'transformer.layers.*.self_attn.q_proj.weight',
+          'model.layers.*.self_attn.q_proj.weight',
           (None, 'model', None),
       ),
       'layers.*.attn.v_proj.w': (
-          'transformer.layers.*.self_attn.v_proj.weight',
+          'model.layers.*.self_attn.v_proj.weight',
           (None, 'model', None),
       ),
-      'final_norm.w': ('transformer.norm.scale', (None,)),
+      'final_norm.w': ('model.norm.scale', (None,)),
   }
 
 

tunix/generate/sglang_jax_sampler.py

@@ -13,58 +13,58 @@ def _to_sglang_jax_mappings() -> Dict[str, MappingEntry]:
   return {
       'lm_head.w': ('lm_head.embedding', (None, 'model')),
       'embedder.input_embedding': (
-          'transformer.embed_tokens.embedding',
+          'model.embed_tokens.embedding',
           ('model', None),
       ),
       'layers.*.input_layernorm.w': (
-          'transformer.layers.*.input_layernorm.scale',
+          'model.layers.*.input_layernorm.scale',
           (None,),
       ),
       'layers.*.mlp.down_proj.kernel': (
-          'transformer.layers.*.mlp.down_proj.weight',
+          'model.layers.*.mlp.down_proj.weight',
           ('model', None),
       ),
       'layers.*.mlp.gate_proj.kernel': (
-          'transformer.layers.*.mlp.gate_proj.weight',
+          'model.layers.*.mlp.gate_proj.weight',
           (None, 'model'),
       ),
       'layers.*.mlp.up_proj.kernel': (
-          'transformer.layers.*.mlp.up_proj.weight',
+          'model.layers.*.mlp.up_proj.weight',
           (None, 'model'),
       ),
       'layers.*.post_attention_layernorm.w': (
-          'transformer.layers.*.post_attention_layernorm.scale',
+          'model.layers.*.post_attention_layernorm.scale',
           (None,),
       ),
       'layers.*.attn.k_proj.w': (
-          'transformer.layers.*.self_attn.k_proj.weight',
+          'model.layers.*.self_attn.k_proj.weight',
           (None, 'model', None),
       ),
       'layers.*.attn.k_bias': (
-          'transformer.layers.*.self_attn.k_proj.bias',
+          'model.layers.*.self_attn.k_proj.bias',
           (None,),
       ),
       'layers.*.attn.o_proj.w': (
-          'transformer.layers.*.self_attn.o_proj.weight',
+          'model.layers.*.self_attn.o_proj.weight',
           ('model', None, None),
       ),
       'layers.*.attn.q_proj.w': (
-          'transformer.layers.*.self_attn.q_proj.weight',
+          'model.layers.*.self_attn.q_proj.weight',
           (None, 'model', None),
       ),
       'layers.*.attn.q_bias': (
-          'transformer.layers.*.self_attn.q_proj.bias',
+          'model.layers.*.self_attn.q_proj.bias',
           (None,),
       ),
       'layers.*.attn.v_proj.w': (
-          'transformer.layers.*.self_attn.v_proj.weight',
+          'model.layers.*.self_attn.v_proj.weight',
           (None, 'model', None),
       ),
       'layers.*.attn.v_bias': (
-          'transformer.layers.*.self_attn.v_proj.bias',
+          'model.layers.*.self_attn.v_proj.bias',
           (None,),
       ),
-      'final_norm.w': ('transformer.norm.scale', (None,)),
+      'final_norm.w': ('model.norm.scale', (None,)),
   }
 
 

tunix/models/llama3/mapping_sglang_jax.py

@@ -26,6 +26,11 @@
 import jax.numpy as jnp
 import safetensors.flax as safetensors
 
+# DO NOT CHNAGE THIS IMPORT. This is used in both oss and GOOGLE_INTERNAL_PACKAGE_PATH.
+from tunix.oss import utils
+
+load_file_from_gcs = utils.load_file_from_gcs
+
 
 def torch_key_to_jax_key(mapping, source_key):
   """Convert torch key to jax key using the provided mapping."""
@@ -78,6 +83,9 @@ def load_and_create_model(
   Returns:
       Model instance with loaded weights
   """
+
+  file_dir = load_file_from_gcs(file_dir)
+
   files = list(epath.Path(file_dir).expanduser().glob("*.safetensors"))
 
   if not files:

tunix/models/qwen2/mapping_sglang_jax.py

@@ -15,6 +15,8 @@
 
 import os
 
+import fsspec
+
 
 def pathways_available() -> bool:
   if "proxy" not in os.getenv("JAX_PLATFORMS", ""):
@@ -25,3 +27,30 @@ def pathways_available() -> bool:
     return True
   except ImportError:
     return False
+
+
+def load_file_from_gcs(file_dir: str, target_dir: str = None) -> str:
+  """Load file from GCS."""
+  if file_dir.startswith("/"):
+    return file_dir
+
+  if not file_dir.startswith("gs://"):
+    raise ValueError(f"Invalid GCS path: {file_dir}")
+
+  _, prefix = file_dir[5:].split("/", 1)
+  try:
+    import tempfile  # pylint: disable=g-import-not-at-top
+
+    if target_dir is None:
+      target_dir = tempfile.gettempdir()
+    local_dir = os.path.join(target_dir, prefix)
+
+    fsspec_fs = fsspec.filesystem("gs")
+    fsspec_fs.get(file_dir, local_dir, recursive=True)
+
+    return local_dir
+  except ImportError as e:
+    raise ImportError(
+        "Please install google-cloud-storage to load model from GCS."
+    ) from e
+