diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index 091a99fb5f..974b513021 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/src/maxtext/configs/post_train/rl.yml @@ -115,6 +115,7 @@ num_batches: 4 # and/or async sampling and training. train_micro_batch_size: -1 rollout_micro_batch_size: -1 +max_seq_token_per_tpu: 0 # > 0 packs sequences into rows of this many tokens for training; 0 = one padded row per sequence # Keep `num_test_batches` low so that evaluation runs quickly. It can be # increased to a max. of 330 (if batch size is 4). num_test_batches: 5 # 200 @@ -176,6 +177,7 @@ enable_dp_attention: false # Performance tuning for samplers max_num_batched_tokens: null max_num_seqs: null +vllm_block_size: null # KV-cache page size; null lets the backend choose # If true, enables asynchronous scheduling in vLLM for faster generation async_scheduling: true # stop generation when any of these strings is generated diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index f6e6782d84..eb1034fd19 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2659,6 +2659,11 @@ class VLLM(BaseModel): async_scheduling: bool = Field(False, description="Enable asynchronous scheduling in vLLM.") max_num_batched_tokens: Optional[int] = Field(None, description="Max number of batched tokens in vLLM.") max_num_seqs: Optional[int] = Field(None, description="Max number of sequences in vLLM.") + vllm_block_size: Optional[int] = Field( + None, + gt=0, + description="KV-cache block (page) size for vLLM. None lets the backend pick it from the engine shape.", + ) stop_strings: Optional[list[str]] = Field(None, description="List of stop strings for vLLM decoding.") vllm_additional_config: dict[str, Any] = Field(default_factory=dict, description="Additional vLLM config options.") vllm_hf_overrides: dict[str, Any] = Field( @@ -2767,6 +2772,15 @@ class RLDataset(BaseModel): train_fraction: float = Field(1.0, gt=0.0, le=1.0, description="Fraction of the dataset to be used for training.") train_micro_batch_size: int = Field(-1, description="Micro batch size for training.") rollout_micro_batch_size: int = Field(-1, description="Micro batch size for rollout.") + max_seq_token_per_tpu: int = Field( + 0, + ge=0, + description=( + "Token budget per packed training row (Tunix `max_seq_token_per_tpu`). When > 0, rollout sequences are packed " + "into rows of this many tokens for the actor/reference passes instead of one padded row per sequence; a maximal " + "sequence (max_prefill_predict_length + generation length) must fit in one row. 0 disables packing." + ), + ) dataset_processor_path: str = Field( "", description=( diff --git a/src/maxtext/integration/tunix/tunix_adapter.py b/src/maxtext/integration/tunix/tunix_adapter.py index 392e746ef3..20a51bef65 100644 --- a/src/maxtext/integration/tunix/tunix_adapter.py +++ b/src/maxtext/integration/tunix/tunix_adapter.py @@ -109,10 +109,18 @@ def __call__( decoder_segment_ids: Optional[Array] = None, output_hidden_states: bool = False, # ignored forced_routed_experts: Optional[Array] = None, + segment_ids: Optional[Array] = None, ) -> Tuple[Array, None]: """Forward compatible with Tunix Trainers default loss. Returns logits, None. + + `segment_ids` is the name Tunix uses for packed-sequence segment ids: it + forwards them only to models whose call signature has a parameter of that + exact name. They are MaxText's `decoder_segment_ids`; when both are given, + `segment_ids` wins so packed rows keep per-sequence attention isolation. """ + if segment_ids is not None: + decoder_segment_ids = segment_ids if decoder_segment_ids is None and self._pad_id is not None: decoder_segment_ids = (input_tokens != self._pad_id).astype(jnp.int32) logits = self.base( diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 6a6a98e9e9..276e3830db 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -436,6 +436,18 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments rl_rollout_engine = functools.partial(MaxTextVllmRollout, maxtext_config=trainer_config) + rollout_vllm_kwargs = { + "hf_overrides": trainer_config.vllm_hf_overrides, + "enable_expert_parallel": sampler_config.enable_expert_parallel, + "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config), + # Ensures vLLM model initializes with correct dtype (not float32 default) + "dtype": trainer_config.weight_dtype.value, + } + if trainer_config.vllm_block_size is not None: + # Pin the KV-cache page size; left unset, the backend derives it from the + # engine shape, so unrelated engine changes can move it. + rollout_vllm_kwargs["block_size"] = trainer_config.vllm_block_size + cluster_config = rl_cluster_lib.ClusterConfig( role_to_mesh={ rl_cluster_lib.Role.ACTOR: actor_mesh, @@ -456,6 +468,7 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments mini_batch_size=trainer_config.batch_size, train_micro_batch_size=train_micro_batch_size, rollout_micro_batch_size=rollout_micro_batch_size, + max_seq_token_per_tpu=trainer_config.max_seq_token_per_tpu or None, metrics_logging_options=metrics_logging_options, profiler_options=profiler_options, checkpoint_root_directory=checkpoint_dir, @@ -484,13 +497,7 @@ def create_rl_components( # pylint: disable=too-many-positional-arguments rollout_vllm_async_scheduling=trainer_config.async_scheduling, rollout_vllm_server_mode=trainer_config.rl.use_agentic_rollout, rollout_vllm_reshard_chunk_size=trainer_config.rl.reshard_chunk_size, - rollout_vllm_kwargs={ - "hf_overrides": trainer_config.vllm_hf_overrides, - "enable_expert_parallel": sampler_config.enable_expert_parallel, - "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config), - # Ensures vLLM model initializes with correct dtype (not float32 default) - "dtype": trainer_config.weight_dtype.value, - }, + rollout_vllm_kwargs=rollout_vllm_kwargs, rollout_vllm_sampling_kwargs={ "stop": trainer_config.stop_strings, "detokenize": trainer_config.stop_strings is not None, diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index 41daec72e9..a641df641b 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/tests/post_training/unit/train_rl_test.py @@ -53,6 +53,17 @@ def test_rl_config_includes_shared_decoder_defaults(self): self.assertFalse(config.enable_mhc_lite) self.assertFalse(config.enable_prefix_caching) + def test_rl_config_packing_and_block_size_defaults(self): + """Packing is off and the vLLM page size is backend-chosen unless set.""" + config = types.RLConfig(model_name="gemma4-26b") + + self.assertEqual(config.max_seq_token_per_tpu, 0) + self.assertIsNone(config.vllm_block_size) + + config = types.RLConfig(model_name="gemma4-26b", max_seq_token_per_tpu=12288, vllm_block_size=128) + self.assertEqual(config.max_seq_token_per_tpu, 12288) + self.assertEqual(config.vllm_block_size, 128) + @pytest.mark.cpu_only def test_rollout_prefix_caching_respects_config_for_attention_model(self): config = SimpleNamespace( diff --git a/tests/post_training/unit/tunix_adapter_test.py b/tests/post_training/unit/tunix_adapter_test.py index ab0679e255..2f91c57fff 100644 --- a/tests/post_training/unit/tunix_adapter_test.py +++ b/tests/post_training/unit/tunix_adapter_test.py @@ -116,6 +116,32 @@ def test_passes_through_explicit_segment_ids_unchanged(self): np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(explicit_seg)) + def test_segment_ids_alias_maps_to_decoder_segment_ids(self): + """Tunix passes packed segment ids under the name `segment_ids`; the + adapter must forward them as MaxText's `decoder_segment_ids` instead of + synthesizing a pad mask.""" + adapter = TunixMaxTextAdapter(base_model=self.base, pad_id=99) + + input_tokens = jnp.array([[10, 11, 12, 20, 21]], dtype=jnp.int32) + positions = jnp.array([[0, 1, 2, 0, 1]], dtype=jnp.int32) + packed_seg = jnp.array([[1, 1, 1, 2, 2]], dtype=jnp.int32) + + adapter(input_tokens, positions, None, None, segment_ids=packed_seg) + + np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(packed_seg)) + + def test_segment_ids_alias_takes_precedence_over_decoder_segment_ids(self): + adapter = TunixMaxTextAdapter(base_model=self.base, pad_id=99) + + input_tokens = jnp.array([[10, 11, 12, 20, 21]], dtype=jnp.int32) + positions = jnp.array([[0, 1, 2, 0, 1]], dtype=jnp.int32) + packed_seg = jnp.array([[1, 1, 1, 2, 2]], dtype=jnp.int32) + other_seg = jnp.array([[7, 7, 7, 7, 7]], dtype=jnp.int32) + + adapter(input_tokens, positions, None, None, decoder_segment_ids=other_seg, segment_ids=packed_seg) + + np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(packed_seg)) + def test_forwards_forced_routed_experts(self): """The stub captures the kwarg but nothing asserted on it, so deleting the