Merge pull request #3040 from AI-Hypercomputer:jimmytsai/fix-ga-in-sft-trainer

PiperOrigin-RevId: 865712234

Author: Google-ML-Automation

src/MaxText/configs/sft-vision-chartqa.yml

@@ -15,6 +15,7 @@
 base_config: "base.yml"
 
 use_sft: True
+use_tunix_gradient_accumulation: True
 use_multimodal: True
 # For vision, the prompt contains image, we only train on completion tokens
 sft_train_on_completion_only: True

src/MaxText/configs/sft-vision-slidevqa.yml

@@ -15,6 +15,7 @@
 base_config: "base.yml"
 
 use_sft: True
+use_tunix_gradient_accumulation: True
 use_multimodal: True
 # For vision, the prompt contains image, we only train on completion tokens
 sft_train_on_completion_only: True

src/MaxText/configs/sft.yml

@@ -15,6 +15,7 @@
 base_config: "base.yml"
 
 use_sft: True
+use_tunix_gradient_accumulation: True
 # sft_train_on_completion_only=False trains on both prompt and completion tokens; trains only on completion tokens otherwise
 sft_train_on_completion_only: True
 packing: True

src/MaxText/configs/types.py

@@ -1056,6 +1056,10 @@ class Optimizer(BaseModel):
   gradient_accumulation_steps: PositiveInt = Field(
       1, description="Number of steps to accumulate gradients before updating."
   )
+  use_tunix_gradient_accumulation: bool = Field(
+      False,
+      description="Whether to use the Tunix implementation for gradient accumulation.",
+  )
   gradient_clipping_threshold: NonNegativeFloat = Field(
       1.0, description="The threshold for gradient clipping. 0 disables clipping."
   )

src/MaxText/input_pipeline/_hf_data_processing.py

@@ -54,7 +54,13 @@ def vision_sft_preprocessing_pipeline(
   """pipeline for multimodal SFT with HF dataset"""
 
   assert len(text_columns) == 2, f"Need two text_columns for query and response, received {text_columns=}"
-  batch_size = global_batch_size // jax.process_count()
+  # Tunix GA requires per-micro-batch slicing at the data level,
+  # whereas Native GA processes the full batch and splits it internally.
+  if config.use_tunix_gradient_accumulation:
+    batch_size = global_batch_size // jax.process_count() // config.gradient_accumulation_steps
+  else:
+    batch_size = global_batch_size // jax.process_count()
+
   if config.enable_data_shuffling:
     dataset = dataset.shuffle(seed=config.data_shuffle_seed)
 
@@ -195,13 +201,21 @@ def preprocessing_pipeline(
     generate_padding_batch=False,
     use_dpo=None,
     use_sft=None,
+    use_tunix_gradient_accumulation=False,
+    num_microbatches=1,
     sft_train_on_completion_only=True,
     grain_worker_count=1,  # only support 0 or 1
     max_segments_per_seq=None,
 ):
   """pipeline for preprocessing HF dataset"""
 
   assert global_batch_size % global_mesh.size == 0, "Batch size should be divisible by number of global devices."
+  # Tunix GA requires per-micro-batch slicing at the data level,
+  # whereas Native GA processes the full batch and splits it internally.
+  if use_tunix_gradient_accumulation:
+    batch_size = global_batch_size // jax.process_count() // num_microbatches
+  else:
+    batch_size = global_batch_size // jax.process_count()
 
   if shuffle:
     dataset = dataset.shuffle(seed=data_shuffle_seed)
@@ -303,15 +317,15 @@ def lists2array(x):
       max_segments = None
     operations.append(
         grain.experimental.PackAndBatchOperation(
-            batch_size=global_batch_size // jax.process_count(),
+            batch_size=batch_size,
             length_struct=length_struct,
             max_sequences_per_bin=max_segments,
         )
     )
     operations.append(_input_pipeline_utils.ReformatPacking(data_column_names))
   else:
     operations.append(_input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id))
-    operations.append(grain.Batch(batch_size=global_batch_size // jax.process_count(), drop_remainder=drop_remainder))
+    operations.append(grain.Batch(batch_size=batch_size, drop_remainder=drop_remainder))
 
   if shift and not use_dpo:
     operations.append(_input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1))
@@ -390,6 +404,8 @@ def make_hf_train_iterator(
         generate_padding_batch=config.generate_padding_batch_train,
         use_dpo=config.use_dpo,
         use_sft=config.use_sft,
+        use_tunix_gradient_accumulation=config.use_tunix_gradient_accumulation,
+        num_microbatches=config.gradient_accumulation_steps,
         sft_train_on_completion_only=config.sft_train_on_completion_only,
         chat_template_path=config.chat_template_path,
         max_segments_per_seq=config.max_segments_per_seq,
@@ -443,6 +459,7 @@ def make_hf_eval_iterator(
         generate_padding_batch=config.generate_padding_batch_eval,
         use_dpo=config.use_dpo,
         use_sft=config.use_sft,
+        num_microbatches=config.gradient_accumulation_steps,
         sft_train_on_completion_only=config.sft_train_on_completion_only,
         chat_template_path=config.chat_template_path,
         max_segments_per_seq=config.max_segments_per_seq,

src/MaxText/sft_trainer.py

@@ -165,7 +165,7 @@ def main(argv: Sequence[str]) -> None:
     os.environ["LIBTPU_INIT_ARGS"] = (
         os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
     )
-  config = pyconfig.initialize(argv)
+  config = pyconfig.initialize(argv, use_tunix_gradient_accumulation=False)
   jax.config.update("jax_use_shardy_partitioner", config.shardy)
   max_utils.print_system_information()
   train_utils.validate_train_config(config)

src/MaxText/train.py

@@ -179,9 +179,13 @@ def loss_fn(model, config, data, dropout_rng, params, is_train=True):
   # Zero1+GA to reduce communication overhead.
   # EPS was used to avoid division by zero, but it's not needed when gradient
   # accumulation is enabled since there's no division.
-  if config.gradient_accumulation_steps > 1:
+  if config.gradient_accumulation_steps > 1 and not config.use_tunix_gradient_accumulation:
     loss = total_loss
   else:
+    # When using Tunix gradient accumulation, we revert to standard normalization.
+    # Unlike the manual accumulation path above, Tunix (via optax.MultiSteps) expects
+    # a normalized loss for each step. It handles the accumulation state
+    # updates and scaling internally.
     loss = total_loss / (total_weights + EPS)
 
   # Calculate and Add MTP Loss