Pack Megatron SFT batches with shared prefix trees

.github/workflows/prek.yml

@@ -227,6 +227,7 @@ jobs:
             tests/unit/test_trainer_rank_weird_shapes.py \
             tests/unit/test_trainer_rank_split.py \
             tests/acceptance/trainer_rank_planner \
+            tests/integration/megatron/test_sft_packing.py::test_sft_packing_preserves_training_targets \
             tests/integration/megatron/model_support/test_dispatcher_graph_retention.py \
             tests/integration/megatron/gdn_shared_prefix/test_gdn_planner_runtime_model.py \
             tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_layout_distributed.py::test_distributed_gdn_cp_layout_all_to_all_roundtrips \

docs/fundamentals/sft-training.mdx

@@ -123,6 +123,28 @@ Set `assistant_turns="last"` to calculate loss only on the final assistant
 message in each JSONL row. Earlier assistant messages remain available as
 context. If omitted, the helper trains on all assistant turns.
 
+## Packed training with Megatron
+
+Megatron automatically packs the tokenized examples in each SFT optimizer batch
+into prefix-tree rows. Identical unsupervised prefixes can share computation;
+assistant-turn selection, prediction targets, and branch-local context are
+preserved. No caller-side packing is required, including when the serverless
+deployment uses Megatron.
+
+`batch_size` still controls examples per optimizer update. The trainer's
+`packed_sequence_length` limits tokens per packed row, independently of the
+per-example context limit. A batch can produce several rows whose gradients are
+accumulated into one update. Packing never combines separate optimizer batches
+or advances the learning-rate schedule for individual rows. Examples too long
+for a packed row are not silently truncated by the packer.
+
+Larger example batches can make better use of packing, but change the optimization
+batch size; ART does not increase it automatically to fill a row. Rows retain
+their actual lengths rather than padding every SFT forward to the capacity.
+
+SFT preparation runs on the CPU before trainer-rank fanout. It shares the packing
+planner with RL, not RL's queue, autotuner, or packing-lookahead orchestration.
+
 ## Distillation
 
 Distillation trains a smaller model on completions from a larger teacher model. Generate responses from the teacher, wrap them as trajectories, and fine-tune:

scripts/ci/trainer-rank-gpu-tests.sh

@@ -20,6 +20,13 @@ test -x "${runtime_python}"
   tests/integration/megatron/lora/test_dynamic_lora_slots.py::test_trainer_rank_custom_parameter_reduction_oracle \
   'tests/integration/megatron/lora/test_dynamic_lora_slots.py::test_trainer_rank_tp_head_backward_matches_unsharded_oracle[2]'
 
+# Keep SFT distributed state and compiler workarounds in separate test processes.
+"${runtime_python}" -m pytest --tb=short \
+  tests/integration/megatron/test_sft_packing.py::test_sft_packing_loss_and_gradients
+
+"${runtime_python}" -m pytest --tb=short \
+  tests/integration/megatron/test_shared_expert_stream_handoff.py::test_compiled_shared_expert_handoff
+
 ART_MEGATRON_CONTEXT_PARALLEL_SIZE=2 \
   "${runtime_python}" -m torch.distributed.run --standalone --nproc-per-node=2 \
     dev/trainer_rank_check.py \

src/art/megatron/compile_workarounds.py

@@ -105,6 +105,15 @@ def _install_te_triton_mask_map_workaround() -> None:
         _disable_attr(permutation, name)
 
 
+def _install_shared_expert_handoff_workaround() -> None:
+    from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+
+    # AOT drops standalone CUDA wait_stream graphs. Keep both ownership
+    # handoffs eager, but leave shared-expert math compiled on its side stream.
+    for name in ("pre_forward_comm", "get_output"):
+        _disable_attr(SharedExpertMLP, name)
+
+
 def install_torch_compile_workarounds(
     config: CompileWorkaroundConfig | None = None,
 ) -> None:
@@ -139,6 +148,8 @@ def _sync_dealloc_fake(
 
     if "context_parallel_attention" in flags:
         _install_context_parallel_attention_workaround()
+    if "shared_expert_stream_handoffs" in flags:
+        _install_shared_expert_handoff_workaround()
     if _SELF_ATTN_LINEAR_PROJ_REDUCE_SCATTER_WORKAROUND_FLAG in flags:
         _install_self_attn_linear_proj_reduce_scatter_workaround()
     if "moe_postprocess" in flags:

src/art/megatron/model_support/handlers/default_dense.py

@@ -32,6 +32,8 @@ def _compile_workaround_flags_for_provider(
         # HybridEP owns native communication, dynamic routing metadata, and
         # side-stream lifetimes. Keep only Megatron's thin flex wrapper eager.
         flags = (*flags, "flex_token_dispatch_combine")
+    if bool(getattr(provider, "moe_shared_expert_overlap", False)):
+        flags = (*flags, "shared_expert_stream_handoffs")
     if (
         bool(getattr(provider, "sequence_parallel", False))
         and int(getattr(provider, "tensor_model_parallel_size", 1) or 1) > 1

src/art/megatron/model_support/handlers/nemotron_h.py

@@ -573,6 +573,8 @@ def configure_provider_for_runtime(self, provider: Any) -> None:
         provider.mtp_num_layers = None
         provider.mtp_hybrid_override_pattern = None
         provider.mtp_loss_scaling_factor = None
+        # Router bias belongs to the frozen base model, not the exported LoRA.
+        provider.moe_router_bias_update_rate = 0.0
         _configure_moe_padding(provider)
         provider.use_mamba_mem_eff_path = True
 

src/art/megatron/runtime/data_plane.py

@@ -38,17 +38,20 @@ def close(self) -> None:
             self._mapped = None
 
 
-class SFTBatchData(BaseModel):
-    """Typed in-memory SFT payload sent directly to warm trainer actors."""
-
+class _SFTBatchStats(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid", frozen=True)
 
-    trajectory_tensors: tuple[dict[str, Any], ...]
     learning_rate: float
     num_trajectories: int
     num_tokens: int
     num_trainable_tokens: int
 
+
+class SFTBatchData(_SFTBatchStats):
+    """One tokenized optimizer batch, before trainer-specific packing."""
+
+    trajectory_tensors: tuple[dict[str, Any], ...]
+
     @model_validator(mode="after")
     def _validate_trajectories(self) -> "SFTBatchData":
         if not self.trajectory_tensors:
@@ -63,6 +66,26 @@ def _validate_trajectories(self) -> "SFTBatchData":
         return self
 
 
+class PackedSFTBatchData(_SFTBatchStats):
+    """Packed microbatches belonging to exactly one optimizer update."""
+
+    rows: tuple[dict[str, Any], ...]
+
+
+def pack_sft_batches(
+    batches: tuple[SFTBatchData, ...], *, seq_len: int
+) -> tuple[PackedSFTBatchData, ...]:
+    from art.preprocessing.sft import pack_sft_batch
+
+    return tuple(
+        PackedSFTBatchData(
+            **batch.model_dump(exclude={"trajectory_tensors"}),
+            rows=pack_sft_batch(batch.trajectory_tensors, seq_len=seq_len),
+        )
+        for batch in batches
+    )
+
+
 def validate_packed_batch(batch: InMemoryPackedBatch) -> None:
     tokens = batch.tensors["tokens"]
     shape = tuple(int(size) for size in tokens.shape)

src/art/megatron/runtime/executor.py

@@ -10,7 +10,7 @@
 from art.utils.safetensors import PreparedSafetensors, SafetensorsLayout
 
 from ..tensor_snapshot import PinnedCpuSnapshotStager
-from .data_plane import InMemoryPackedBatch, SFTBatchData, validate_packed_batch
+from .data_plane import InMemoryPackedBatch, PackedSFTBatchData, validate_packed_batch
 from .publication import (
     TrainerPublicationFailed,
     TrainerPublicationSucceeded,
@@ -82,7 +82,7 @@ def execute(
     def execute_sft(
         self,
         job: SFTJobSpec,
-        batches: tuple[SFTBatchData, ...],
+        batches: tuple[PackedSFTBatchData, ...],
         sink: EventSink,
         cancelled: Event,
     ) -> dict[str, float]:

src/art/megatron/runtime/monarch.py

@@ -23,7 +23,12 @@
 from art.utils.cache_dirs import configure_model_cache_env
 from art.utils.lifecycle import cleanup_after_failure, consume_future_exception
 
-from .data_plane import InMemoryPackedBatch, SFTBatchData
+from .data_plane import (
+    InMemoryPackedBatch,
+    PackedSFTBatchData,
+    SFTBatchData,
+    pack_sft_batches,
+)
 from .publication import (
     TRAINER_PUBLICATION_EVENT_ADAPTER,
     TrainerPublicationEvent,
@@ -583,7 +588,7 @@ def execute(
     def execute_sft(
         self,
         job_json: str,
-        batches: tuple[SFTBatchData, ...],
+        batches: tuple[PackedSFTBatchData, ...],
         event_port: Port[dict[str, Any]],
     ) -> dict[str, Any]:
         try:
@@ -1122,11 +1127,19 @@ async def train(
     async def train_sft(
         self, job: SFTJobSpec, batches: tuple[SFTBatchData, ...]
     ) -> AsyncIterator[TrainEvent]:
+        async def dispatch(port: Port[dict[str, Any]]) -> Any:
+            packed = await asyncio.to_thread(
+                pack_sft_batches,
+                batches,
+                seq_len=self.runtime_spec.packed_sequence_length,
+            )
+            return await self._actors.execute_sft.call(
+                job.model_dump_json(), packed, port
+            )
+
         async for event in self._train(
             job,
-            lambda port: self._actors.execute_sft.call(
-                job.model_dump_json(), batches, port
-            ),
+            dispatch,
             lambda: self._validate_sft(job, batches),
         ):
             yield event

src/art/megatron/train.py

@@ -69,7 +69,7 @@
     build_moe_routing_replay_bundle_from_packed_tensors,
     prepare_moe_routing_replay_boundaries,
 )
-from art.megatron.runtime.data_plane import SFTBatchData
+from art.megatron.runtime.data_plane import PackedSFTBatchData
 from art.megatron.runtime.specs import (
     PackedTokenScore,
     ResidentLoraExport,
@@ -116,6 +116,7 @@
     select_indexed_inputs,
     select_micro_inputs,
     select_sft_micro_inputs,
+    sft_global_microbatch_count,
 )
 from art.megatron.training.model_chunks import (
     ModelChunks,
@@ -828,7 +829,7 @@ def execute_megatron_rl_job(
 def execute_megatron_sft_job(
     runtime: TrainingRuntime,
     job: SFTJobSpec,
-    batches: tuple[SFTBatchData, ...],
+    batches: tuple[PackedSFTBatchData, ...],
     *,
     progress_sink: Callable[[int, int, dict[str, float]], None],
     adapter_ready_sink: Callable[[], None],
@@ -847,10 +848,6 @@ def execute_megatron_sft_job(
     try:
         configure_moe_routing_replay(runtime)
         adapter_dtypes = _prepare_rl_training_state(runtime, job)
-        grad_accumulation_sequences = int(job.config.batch_size)
-        grad_accumulation_sequences = resolve_global_grad_accumulation_sequences(
-            grad_accumulation_sequences
-        )
         assert runtime.optimizer is not None
         runtime.optimizer.config.clip_grad = job.max_grad_norm
         for param_group in runtime.optimizer.param_groups:
@@ -863,20 +860,18 @@ def execute_megatron_sft_job(
 
                 raise TrainingCancelledError("SFT job was cancelled")
             started = time.perf_counter()
-            trajectory_tensors = list(batch.trajectory_tensors)
+            trajectory_tensors = list(batch.rows)
+            grad_accumulation_sequences = sft_global_microbatch_count(
+                len(trajectory_tensors), provider=runtime.provider
+            )
             template = _clone_sft_tensors(trajectory_tensors[0])
             zero_template = _zero_contribution_sft_inputs(template)
-            # Scheduling uses run-global sample IDs while each payload only owns one
-            # batch. Prefix aliases place this window in global index space without
-            # copying tensors, then selected IDs are rebased for local lookup.
-            sample_offset = batch_index * grad_accumulation_sequences
-            scheduled_tensors = [
-                trajectory_tensors[0]
-            ] * sample_offset + trajectory_tensors
+            # Packed rows, not source examples, are schedule microbatches. Each
+            # payload still owns exactly one optimizer update and learning rate.
             hybridep_token_counts = (
                 build_sft_hybridep_token_counts(
-                    trajectory_tensors=scheduled_tensors,
-                    step_index=batch_index,
+                    trajectory_tensors=trajectory_tensors,
+                    step_index=0,
                     global_grad_accumulation_sequences=(grad_accumulation_sequences),
                     topology=topology,
                     provider=runtime.provider,
@@ -894,14 +889,11 @@ def execute_megatron_sft_job(
                 required_capacity=max(hybridep_token_counts or (), default=0),
             )
             scheduled_indices = build_micro_sample_indices(
-                step_index=batch_index,
-                num_sequences=len(scheduled_tensors),
+                step_index=0,
+                num_sequences=len(trajectory_tensors),
                 global_grad_accumulation_sequences=grad_accumulation_sequences,
             )
-            micro_indices = [
-                None if index is None else index - sample_offset
-                for index in scheduled_indices
-            ]
+            micro_indices = scheduled_indices
             step_result = run_megatron_sft_step(
                 model_chunks=runtime.model,
                 provider=runtime.provider,
@@ -919,13 +911,26 @@ def execute_megatron_sft_job(
                     runtime.optimizer_snapshot_barrier.wait_before_mutation
                 ),
             )
+            _validate_train_step_result_finite(runtime, step_result)
             elapsed = time.perf_counter() - started
             final_metrics = {
                 "loss/train": float(step_result.reduced_loss.item()),
                 "loss/learning_rate": batch.learning_rate,
                 "loss/grad_norm": float(step_result.grad_norm),
                 "throughput/train_executed_tok_equiv_per_s": (
-                    batch.num_tokens / elapsed if elapsed else 0.0
+                    step_result.workload.executed_token_equivalents / elapsed
+                    if elapsed
+                    else 0.0
+                ),
+                "throughput/train_logical_tok_per_s": batch.num_tokens / elapsed
+                if elapsed
+                else 0.0,
+                "data/gradient_step_nonpadding_logical_tokens": float(batch.num_tokens),
+                "data/gradient_step_loss_bearing_tokens": float(
+                    step_result.workload.loss_bearing_tokens
+                ),
+                "data/gradient_step_executed_token_equivalents": float(
+                    step_result.workload.executed_token_equivalents
                 ),
                 **step_result.pipeline_metrics,
             }