Plumb config.max_segments_per_seq to grain dataset max_sequences_per_bin

gabeweisz · gabeweisz · commit f7971f2e2816 · 2025-12-19T16:47:01.000Z
diff --git a/src/MaxText/configs/a3/llama_3.1_405b/128vm.sh b/src/MaxText/configs/a3/llama_3.1_405b/128vm.sh
@@ -52,5 +52,5 @@ python3 -m MaxText.$EXECUTABLE ${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src
     dcn_fsdp_parallelism=128 \
     ici_fsdp_parallelism=8 \
     base_output_directory=$OUTPUT_PATH \
+    max_segments_per_seq=32 \
     profiler=xplane
-
diff --git a/src/MaxText/configs/base.yml b/src/MaxText/configs/base.yml
@@ -552,9 +552,10 @@ generate_padding_batch_train: False
 generate_padding_batch_eval: False
 # Maximum number of segments that can be packed into a single sequence
 # This needs to be passed to TransformerEngine's DotProductAttention layer for packing
-# NOTE: This parameter is only relevant for GPU packed attention.
-# When using TPU, this parameter will be ignored.
-max_segments_per_seq: 32
+# This also affects packing for grain, since TransformerEngine may crash or cause
+# data corruption if there are more segments packed than specified
+# Set this to something like 32 for GPUs when using TransformerEngine
+max_segments_per_seq: -1
 # Rampup batch size, similar to Megatron-LM, see
 # https://github.com/NVIDIA/Megatron-LM/blob/2a01637aa54ccdaf7ea9afc1f1b80f58c53d7f3c/megatron/core/num_microbatches_calculator.py#L233-L237
 # The ramp-up proceeds in stages from `per_device_batch_size_start` up to
diff --git a/src/MaxText/configs/gpu_smoke_test.yml b/src/MaxText/configs/gpu_smoke_test.yml
@@ -12,3 +12,4 @@ per_device_batch_size: 2
 max_target_length: 1024
 dataset_type: "synthetic"
 steps: 10
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/llama2_70b.yml b/src/MaxText/configs/models/gpu/llama2_70b.yml
@@ -16,3 +16,4 @@ logits_dot_in_fp32: False
 
 per_device_batch_size: 6
 max_target_length: 4096
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/llama2_7b.yml b/src/MaxText/configs/models/gpu/llama2_7b.yml
@@ -12,4 +12,5 @@ remat_policy: "minimal_with_context"
 use_iota_embed: True
 scan_layers: False
 dataset_type: "synthetic"
-async_checkpointing: False
+async_checkpointing: False
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/llama3.1_405b.yml b/src/MaxText/configs/models/gpu/llama3.1_405b.yml
@@ -14,3 +14,4 @@ async_checkpointing: False
 logits_dot_in_fp32: False
 per_device_batch_size: 1.0
 max_target_length: 4096
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/llama3_70b.yml b/src/MaxText/configs/models/gpu/llama3_70b.yml
@@ -28,3 +28,4 @@ use_iota_embed: True
 dataset_type: "synthetic"
 reuse_example_batch: 1
 enable_checkpointing: False
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/llama3_8b.yml b/src/MaxText/configs/models/gpu/llama3_8b.yml
@@ -28,3 +28,4 @@ use_iota_embed: True
 dataset_type: "synthetic"
 reuse_example_batch: 1
 enable_checkpointing: False
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/mixtral_8x1b.yml b/src/MaxText/configs/models/gpu/mixtral_8x1b.yml
@@ -33,3 +33,4 @@ reuse_example_batch: 1
 enable_checkpointing: False
 megablox: False
 sparse_matmul: False
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/mixtral_8x2b.yml b/src/MaxText/configs/models/gpu/mixtral_8x2b.yml
@@ -33,3 +33,4 @@ reuse_example_batch: 1
 enable_checkpointing: False
 megablox: False
 sparse_matmul: False
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/models/gpu/mixtral_8x7b.yml b/src/MaxText/configs/models/gpu/mixtral_8x7b.yml
@@ -33,3 +33,4 @@ scan_layers: False
 tokenizer_path: "/deps/src/MaxText/assets/tokenizer.mistral-v1"
 profiler: "nsys"
 capacity_factor: 1.0
+max_segments_per_seq: 32
diff --git a/src/MaxText/configs/types.py b/src/MaxText/configs/types.py
@@ -279,9 +279,7 @@ class Checkpointing(BaseModel):
   save_checkpoint_on_completion: bool = Field(
       True, description="If True, saves a final checkpoint upon training completion."
   )
-  enable_continuous_checkpointing: bool = Field(
-      False, description="If True, enables continuous checkpointing."
-  )
+  enable_continuous_checkpointing: bool = Field(False, description="If True, enables continuous checkpointing.")
 
 
 class OrbaxStorage(BaseModel):
@@ -463,9 +461,7 @@ class Attention(BaseModel):
   ragged_block_size: int = Field(256, description="Block size for ragged attention.")
   enable_padding_causal_mask: bool = Field(True, description="Temporary flag for TE padding.")
   use_tokamax_splash: bool = Field(False, description="Whether to use tokamax splash attention.")
-  use_jax_splash: bool = Field(
-      False, description="Whether to use jax splash attention."
-  )
+  use_jax_splash: bool = Field(False, description="Whether to use jax splash attention.")
 
 
 class MoBa(BaseModel):
@@ -880,8 +876,8 @@ class DatasetGeneral(BaseModel):
       description="Packing type when using Grain pipeline. 'first_fit' or 'concat_then_split'.",
   )
   max_segments_per_seq: int = Field(
-      32,
-      description="Maximum number of segments that can be packed into a single sequence.",
+      -1,
+      description="Maximum number of segments that can be packed into a single sequence. -1 or None for no limit.",
   )
   num_epoch: int = Field(1, description="Number of epochs to train for.")
   expansion_factor_real_data: float = Field(-1.0, description="Factor for partial data loading on hosts.")
@@ -2065,6 +2061,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError(
             "Ring context parallelism strategy (context_parallel_strategy='ring') is only supported on GPUs."
         )
+    if self.hardware == "gpu" and self.packing and self.attention == "cudnn_flash_te" and self.max_segments_per_seq <= 0:
+      raise ValueError("max_segments_per_seq must be set when using TransformerEngine attention and packing")
     dcn_product = (
         self.dcn_data_parallelism
         * self.dcn_pipeline_parallelism
diff --git a/src/MaxText/input_pipeline/_grain_data_processing.py b/src/MaxText/input_pipeline/_grain_data_processing.py
@@ -242,9 +242,15 @@ def pretrain_preprocessing_pipeline(
 
   if config.packing:
     length_struct = {col: config.max_target_length for col in data_columns}
+    max_segments = config.max_segments_per_seq
+    if max_segments is not None and max_segments <= 0:
+      max_segments = None
     if config.grain_packing_type == "first_fit":
       dataset = grain.experimental.FirstFitPackIterDataset(
-          dataset, length_struct=length_struct, num_packing_bins=batch_size
+          dataset,
+          length_struct=length_struct,
+          num_packing_bins=batch_size,
+          max_sequences_per_bin=max_segments,
       )
     elif config.grain_packing_type == "concat_then_split":
       if config.add_bos and hasattr(tokenizer_model, "bos_id"):
diff --git a/src/MaxText/input_pipeline/_hf_data_processing.py b/src/MaxText/input_pipeline/_hf_data_processing.py
@@ -192,6 +192,7 @@ def preprocessing_pipeline(
     use_sft=None,
     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"""
 
@@ -297,10 +298,14 @@ def lists2array(x):
 
   if packing and not use_dpo:
     length_struct = {col: max_target_length for col in data_column_names}
+    max_segments = max_segments_per_seq
+    if max_segments is not None and max_segments <= 0:
+      max_segments = None
     operations.append(
         grain.experimental.PackAndBatchOperation(
             batch_size=global_batch_size // jax.process_count(),
             length_struct=length_struct,
+            max_sequences_per_bin=max_segments,
         )
     )
     operations.append(_input_pipeline_utils.ReformatPacking(data_column_names))
@@ -386,6 +391,7 @@ def make_hf_train_iterator(
         use_sft=config.use_sft,
         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,
     )
   return train_iter
 
@@ -437,5 +443,6 @@ def make_hf_eval_iterator(
         use_sft=config.use_sft,
         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,
     )
   return eval_iter

Original file line number	Diff line number	Diff line change
`@@ -16,3 +16,4 @@ logits_dot_in_fp32: False`
`16`	`16`
`17`	`17`	`per_device_batch_size: 6`
`18`	`18`	`max_target_length: 4096`
	`19`	`+max_segments_per_seq: 32`