diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 97af92710..6c8a29d72 100644 --- a/.github/workflows/prek.yml +++ b/.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 \ diff --git a/docs/fundamentals/sft-training.mdx b/docs/fundamentals/sft-training.mdx index e49fe475e..45e89b20f 100644 --- a/docs/fundamentals/sft-training.mdx +++ b/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: diff --git a/scripts/ci/trainer-rank-gpu-tests.sh b/scripts/ci/trainer-rank-gpu-tests.sh index 7e4396a38..b496c80b4 100755 --- a/scripts/ci/trainer-rank-gpu-tests.sh +++ b/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 \ diff --git a/src/art/megatron/compile_workarounds.py b/src/art/megatron/compile_workarounds.py index 4c2ec6dd5..0759ace4b 100644 --- a/src/art/megatron/compile_workarounds.py +++ b/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: diff --git a/src/art/megatron/model_support/handlers/default_dense.py b/src/art/megatron/model_support/handlers/default_dense.py index 360a07e2f..57c0e787a 100644 --- a/src/art/megatron/model_support/handlers/default_dense.py +++ b/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 diff --git a/src/art/megatron/model_support/handlers/nemotron_h.py b/src/art/megatron/model_support/handlers/nemotron_h.py index d6d1ce5cd..c10db5d6c 100644 --- a/src/art/megatron/model_support/handlers/nemotron_h.py +++ b/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 diff --git a/src/art/megatron/runtime/data_plane.py b/src/art/megatron/runtime/data_plane.py index d5f83ddd5..9223e2a83 100644 --- a/src/art/megatron/runtime/data_plane.py +++ b/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) diff --git a/src/art/megatron/runtime/executor.py b/src/art/megatron/runtime/executor.py index 77ec83cce..5e9e54e1e 100644 --- a/src/art/megatron/runtime/executor.py +++ b/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]: diff --git a/src/art/megatron/runtime/monarch.py b/src/art/megatron/runtime/monarch.py index 29f3db7ca..c07f5e0b4 100644 --- a/src/art/megatron/runtime/monarch.py +++ b/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 diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py index 2311d8050..cf7d9452a 100644 --- a/src/art/megatron/train.py +++ b/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, } diff --git a/src/art/megatron/training/microbatches.py b/src/art/megatron/training/microbatches.py index ad5a886c8..1da109075 100644 --- a/src/art/megatron/training/microbatches.py +++ b/src/art/megatron/training/microbatches.py @@ -317,6 +317,20 @@ def rank_counts(sample_index: int | None) -> tuple[int, ...]: ] +def sft_global_microbatch_count(num_rows: int, *, provider: Any) -> int: + """Round only the execution schedule, never the optimizer batch or loss count.""" + dp = ps.get_data_parallel_world_size() + local = (num_rows + dp - 1) // dp + if (ps.get_virtual_pipeline_model_parallel_world_size() or 1) > 1: + pp = ps.get_pipeline_model_parallel_world_size() + group = int(getattr(provider, "microbatch_group_size_per_vp_stage", 0) or pp) + local = max(local, group) + remainder = local % group + if 0 < remainder < pp: + local += pp - remainder + return local * dp + + def build_sft_hybridep_token_counts( *, trajectory_tensors: list[dict[str, torch.Tensor]], @@ -817,7 +831,11 @@ def _prepare_dense_sft_micro( seq_len = max(int(attention_mask.sum().item()), 1) input_ids = micro["input_ids"].reshape(-1)[:seq_len].unsqueeze(0).to(device) labels = micro["labels"].reshape(-1)[:seq_len].unsqueeze(0) - position_ids = torch.arange(seq_len, device=device).unsqueeze(0) + position_ids = ( + micro["input_pos"].to(device) + if "input_pos" in micro + else torch.arange(seq_len, device=device).unsqueeze(0) + ) shifted_labels = shift_tensor(labels, -100) loss_mask = shifted_labels != -100 workload = TrainingMicrobatchWorkload( @@ -832,13 +850,25 @@ def _prepare_dense_sft_micro( ) shifted_labels = shifted_labels.to(device) loss_mask = loss_mask.to(device) - return PreparedSFTMicroInputs( - input_ids=input_ids, - position_ids=position_ids, - labels=shifted_labels, - loss_mask=loss_mask, - lm_head_selection=lm_head_selection, - attention_state=_causal_attention_state( + attention_state = ( + create_prefix_tree_state( + group_ids=micro["group_ids"], + parent_ids=micro["parent_ids"], + input_pos=micro["input_pos"], + target_device=device, + sliding_windows=_art_flex_sliding_windows(provider), + build_gdn_execution_spec=bool( + getattr(model_support_handler, "build_gdn_execution_spec", False) + ), + gdn_planner_config=_gdn_planner_config_for_provider( + provider, model_support_handler + ), + model_support_handler=model_support_handler, + attention_head_dim=getattr(provider, "kv_channels", None), + attention_value_head_dim=getattr(provider, "kv_channels", None), + ) + if "group_ids" in micro + else _causal_attention_state( seq_len, device, sliding_windows=_art_flex_sliding_windows(provider), @@ -849,7 +879,15 @@ def _prepare_dense_sft_micro( model_support_handler=model_support_handler, attention_head_dim=getattr(provider, "kv_channels", None), attention_value_head_dim=getattr(provider, "kv_channels", None), - ), + ) + ) + return PreparedSFTMicroInputs( + input_ids=input_ids, + position_ids=position_ids, + labels=shifted_labels, + loss_mask=loss_mask, + lm_head_selection=lm_head_selection, + attention_state=attention_state, local_token_uids=sft_sequence_token_uids(micro, device=device)[ :, : int(input_ids.shape[1]) ], @@ -872,14 +910,21 @@ def _sft_inputs_to_sparse_packed_tensors( parent_ids = torch.full((1, total_tokens), -1, device=device, dtype=torch.long) group_ids[:, :actual_len] = 0 parent_ids[:, :actual_len] = 0 + if "group_ids" in inputs: + group_ids = inputs["group_ids"].to(device) + parent_ids = inputs["parent_ids"].to(device) assistant_mask = (labels != -100).unsqueeze(0).to(device=device, dtype=torch.bool) return PackedTensors( tokens=input_ids.unsqueeze(0).to(device=device, dtype=torch.long), group_ids=group_ids, parent_ids=parent_ids, - input_pos=torch.arange(total_tokens, device=device, dtype=torch.long).unsqueeze( - 0 + input_pos=( + inputs["input_pos"].to(device) + if "input_pos" in inputs + else torch.arange(total_tokens, device=device, dtype=torch.long).unsqueeze( + 0 + ) ), assistant_mask=assistant_mask, logprobs=torch.full( diff --git a/src/art/preprocessing/pack.py b/src/art/preprocessing/pack.py index a4149039a..dc50c5dbf 100644 --- a/src/art/preprocessing/pack.py +++ b/src/art/preprocessing/pack.py @@ -2,7 +2,7 @@ import os import random import time -from typing import Any, Literal, NamedTuple, cast +from typing import Any, Literal, NamedTuple, Protocol, TypeVar, cast import numpy as np import torch @@ -28,6 +28,20 @@ DEFAULT_MIN_PREFIX_TREE_SHARED_SEGMENT_LENGTH = 64 +class PrefixTreeSequence(Protocol): + @property + def token_ids(self) -> tuple[int, ...]: ... + + @property + def shareable_length(self) -> int: ... + + @property + def prompt_id(self) -> int: ... + + +_Sequence = TypeVar("_Sequence", bound=PrefixTreeSequence) + + class PrefixTreePackingStats(TypedDict): logical_tokens: int physical_tokens: int @@ -382,12 +396,13 @@ def prefix_tree_pack( def _prefix_tree_pack_rows( - items: list[_PrefixTreePackItem], + items: list[_Sequence], *, seq_len: int, pack_results: bool, min_shared_segment_length: int, -) -> list[tuple[list[_PrefixTreePackItem], _PrefixTreeRowPlan]]: + rebuild_rows: bool = True, +) -> list[tuple[list[_Sequence], _PrefixTreeRowPlan]]: if not items: return [] if not pack_results: @@ -447,6 +462,9 @@ def _prefix_tree_pack_rows( "Global prefix-tree occupancy disagrees with final bin plan: " f"occupancy={packed_bin.token_count}, plan={occupancy_plan.length}" ) + if not rebuild_rows: + planned_rows.append((row, occupancy_plan)) + continue # Rebuild only after placement so bin-local paths compress without putting # repeated tree construction in the best-fit search. plan = _prefix_tree_row_plan( @@ -669,7 +687,7 @@ def _first_trainable_token_index( def _prefix_tree_row_plan( - row: list[_PrefixTreePackItem], + row: Sequence[PrefixTreeSequence], *, seq_len: int, pack_results: bool, diff --git a/src/art/preprocessing/sft.py b/src/art/preprocessing/sft.py new file mode 100644 index 000000000..ea101f1a1 --- /dev/null +++ b/src/art/preprocessing/sft.py @@ -0,0 +1,79 @@ +"""Loss-preserving SFT materialization over the shared prefix-tree row planner.""" + +from collections.abc import Sequence + +import numpy as np +from pydantic import BaseModel, ConfigDict +import torch + +from .pack import DEFAULT_MIN_PREFIX_TREE_SHARED_SEGMENT_LENGTH, _prefix_tree_pack_rows + + +class _SFTSequence(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) + + token_ids: tuple[int, ...] + labels: np.ndarray + shareable_length: int + prompt_id: int + + +def pack_sft_batch( + inputs: Sequence[dict[str, torch.Tensor]], *, seq_len: int +) -> tuple[dict[str, torch.Tensor], ...]: + """Pack one optimizer batch without filtering, truncating, or changing labels. + + Rows retain natural lengths. Supervised prediction positions are never shared; + labels are still unshifted, as in tokenize_sft_batch's output. + """ + items = [] + for index, tensors in enumerate(inputs): + length = int(tensors["attention_mask"].sum().item()) + tokens = tensors["input_ids"].reshape(-1)[:length].numpy() + labels = tensors["labels"].reshape(-1)[:length].numpy().copy() + # An example's first token has no preceding hidden state to supervise. + if length: + labels[0] = -100 + targets = np.flatnonzero(labels != -100) + items.append( + _SFTSequence( + token_ids=tuple(tokens.tolist()), + labels=labels, + shareable_length=max(int(targets[0]) - 1, 0) + if targets.size + else length, + prompt_id=index, + ) + ) + rows = [] + for sequences, plan in _prefix_tree_pack_rows( + items, + seq_len=seq_len, + pack_results=True, + min_shared_segment_length=DEFAULT_MIN_PREFIX_TREE_SHARED_SEGMENT_LENGTH, + # Rebuilding can change sharing at unequal supervision boundaries and + # exceed the admitted capacity. Keep the already compacted bin geometry. + rebuild_rows=False, + ): + fields = { + name: np.empty((1, plan.length), dtype=np.int64) + for name in ("input_ids", "labels", "group_ids", "parent_ids", "input_pos") + } + for segment in plan.segments: + sequence = sequences[segment.sequence_indices[0]] + source = slice(segment.start, segment.end) + dest = slice(segment.packed_start, segment.packed_start + segment.length) + fields["input_ids"][0, dest] = sequence.token_ids[source] + fields["labels"][0, dest] = sequence.labels[source] + fields["group_ids"][0, dest] = segment.group_id + fields["parent_ids"][0, dest] = segment.parent_id + fields["input_pos"][0, dest] = np.arange(segment.start, segment.end) + rows.append( + { + **{name: torch.from_numpy(value) for name, value in fields.items()}, + "attention_mask": torch.ones( + (1, plan.length), dtype=torch.long, device="cpu" + ), + } + ) + return tuple(rows) diff --git a/tests/integration/megatron/test_sft_packing.py b/tests/integration/megatron/test_sft_packing.py new file mode 100644 index 000000000..50fd5aab7 --- /dev/null +++ b/tests/integration/megatron/test_sft_packing.py @@ -0,0 +1,260 @@ +"""Real SFT loss/gradient parity across packed rows and independent examples. + +Run with pytest on one GPU, or torchrun this file for CP/DP/PP qualification. +ART_SFT_TEST_MODEL optionally selects a pretrained checkpoint instead of the tiny +random Llama fixture; ART_SFT_TEST_SUPPORT_KEY supplies its ART handler key. +""" + +from collections import Counter +import json +import os +from pathlib import Path +import tempfile + +import pytest +import torch + + +def _examples(last_only: bool) -> list[dict[str, torch.Tensor]]: + generator = torch.Generator().manual_seed(123) + prefix = torch.randint(1, 128, (128,), generator=generator) + examples = [] + for index, length in enumerate((193, 217, 241, 205)): + tokens = torch.cat( + (prefix, torch.randint(1, 128, (length - 128,), generator=generator)) + ) + if index == 3: + tokens[:64] = tokens[:64].flip(0) + labels = tokens.clone() + labels[:176] = -100 + if not last_only: + labels[128:144] = tokens[128:144] + examples.append( + { + "input_ids": tokens[None], + "labels": labels[None], + "attention_mask": torch.ones_like(tokens)[None], + } + ) + return examples + + +def test_sft_packing_preserves_training_targets(): + from art.megatron.prefix_tree import parse_prefix_tree + from art.preprocessing.sft import pack_sft_batch + + cases = [_examples(False), _examples(True)] + unequal = [] + for index, (length, first_target) in enumerate( + ((87, 87), (116, 72), (124, 88), (275, 244), (250, 148)) + ): + tokens = torch.arange(length)[None] + tokens[:, 64 if index == 0 else 128 :] += (index + 1) * 1000 + labels = tokens.clone() + labels[:, :first_target] = -100 + unequal.append( + dict( + input_ids=tokens, labels=labels, attention_mask=torch.ones_like(tokens) + ) + ) + cases.append(unequal) + for examples in cases: + expected = Counter( + (tuple(example["input_ids"][0, :index].tolist()), int(label)) + for example in examples + for index, label in enumerate(example["labels"][0]) + if index > 0 and label != -100 + ) + for capacity in (384, 512, 1024): + actual = Counter() + for row in pack_sft_batch(examples, seq_len=capacity): + (tree,) = parse_prefix_tree( + group_ids=row["group_ids"], parent_ids=row["parent_ids"] + ) + segments = {segment.group_id: segment for segment in tree.segments} + tokens, labels = row["input_ids"][0].tolist(), row["labels"][0].tolist() + for segment in tree.segments: + prefix = [ + token + for ancestor in segment.ancestors + for token in tokens[ + segments[ancestor].start : segments[ancestor].end + ] + ] + for position in range( + segment.start, min(segment.end, len(tokens) - 1) + ): + if labels[position + 1] != -100: + actual[ + ( + tuple( + prefix + tokens[segment.start : position + 1] + ), + labels[position + 1], + ) + ] += 1 + assert actual == expected + with pytest.raises(RuntimeError, match="exceeds sequence length"): + pack_sft_batch(examples, seq_len=64) + + +def _run(runtime, inputs): + from art.megatron.train import run_megatron_sft_step + from art.megatron.training.microbatches import ( + _zero_contribution_sft_inputs, + build_micro_sample_indices, + select_sft_micro_inputs, + sft_global_microbatch_count, + ) + + indices = build_micro_sample_indices( + step_index=0, + num_sequences=len(inputs), + global_grad_accumulation_sequences=sft_global_microbatch_count( + len(inputs), provider=runtime.provider + ), + ) + + gradients = {} + + def capture_gradients(): + for index, model in enumerate(runtime.model): + for name, parameter in model.named_parameters(): + if parameter.requires_grad: + gradients[f"{index}.{name}"] = ( + parameter.main_grad.detach().float().cpu().clone() + ) + + result = run_megatron_sft_step( + model_chunks=runtime.model, + provider=runtime.provider, + model_support_handler=runtime.model_support_handler, + optimizer=runtime.optimizer, + learning_rate=0.0, + inputs=select_sft_micro_inputs( + inputs, indices, _zero_contribution_sft_inputs(inputs[0]) + ), + step_index=0, + sample_index=indices, + before_optimizer_step=capture_gradients, + ) + return float(result.reduced_loss), gradients + + +def _mape(candidate, reference, name): + assert torch.isfinite(candidate).all() and torch.isfinite(reference).all() + assert candidate.abs().sum() > 0 and reference.abs().sum() > 0, name + return float((candidate - reference).abs().mean() / reference.abs().mean() * 100) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires Megatron CUDA runtime" +) +def test_sft_packing_loss_and_gradients(): + from transformers import LlamaConfig + + from art.megatron.train import build_training_runtime + from art.preprocessing.sft import pack_sft_batch + + torch.set_num_threads(4) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + os.environ.setdefault("ART_MEGATRON_LORA_RANK", "8") + for axis in ("TENSOR_MODEL", "CONTEXT", "PIPELINE_MODEL", "EXPERT_MODEL"): + os.environ.setdefault(f"ART_MEGATRON_{axis}_PARALLEL_SIZE", "1") + os.environ.setdefault( + "ART_MEGATRON_LORA_TARGET_MODULES", '["q_proj","k_proj","v_proj","o_proj"]' + ) + if "WORLD_SIZE" in os.environ: + torch.distributed.init_process_group("nccl") + else: + os.environ.update( + RANK="0", WORLD_SIZE="1", LOCAL_RANK="0", LOCAL_WORLD_SIZE="1" + ) + torch.distributed.init_process_group( + "nccl", store=torch.distributed.HashStore(), rank=0, world_size=1 + ) + try: + with tempfile.TemporaryDirectory(prefix="art-sft-packing-") as directory: + checkpoint = os.environ.get("ART_SFT_TEST_MODEL") + if checkpoint is None: + LlamaConfig( + architectures=["LlamaForCausalLM"], + hidden_size=256, + intermediate_size=512, + num_hidden_layers=max( + 2, + int( + os.environ.get( + "ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE", "1" + ) + ) + * int( + os.environ.get( + "ART_MEGATRON_VIRTUAL_PIPELINE_MODEL_PARALLEL_SIZE", "1" + ) + ), + ), + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=256, + max_position_embeddings=2048, + ).save_pretrained(directory) + torch.manual_seed(3407) + runtime = build_training_runtime( + model_identifier=checkpoint or directory, + model_initialization="pretrained" if checkpoint else "random", + model_support_key=os.environ.get( + "ART_SFT_TEST_SUPPORT_KEY", "llama3_dense" + ), + print_env=False, + ) + for model in runtime.model: + model.train() + for parameter in model.parameters(): + if parameter.requires_grad: + with torch.no_grad(): + parameter.normal_(std=0.02) + assert runtime.optimizer is not None + runtime.optimizer.reload_model_params() + router_biases = [ + (buffer, buffer.clone()) + for model in runtime.model + for name, buffer in model.named_buffers() + if name.endswith("expert_bias") + ] + rows = [] + for last_only in (False, True): + examples = _examples(last_only) + expected_loss, expected_gradients = _run(runtime, examples) + for capacity in (384, 1024): + packed = list(pack_sft_batch(examples, seq_len=capacity)) + loss, gradients = _run(runtime, packed) + loss_error = abs(loss - expected_loss) / abs(expected_loss) * 100 + assert expected_loss > 0 and loss > 0 and loss_error <= 3.0 + errors = { + name: _mape(gradients[name], reference, name) + for name, reference in expected_gradients.items() + } + assert max(errors.values()) <= 5.0, errors + rows.append( + dict( + last_only=last_only, + capacity=capacity, + packed_rows=len(packed), + loss_mape=loss_error, + max_gradient_mape=max(errors.values()), + ) + ) + for buffer, initial in router_biases: + torch.testing.assert_close(buffer, initial, rtol=0, atol=0) + print(json.dumps(rows), flush=True) + if output := os.environ.get("ART_SFT_TEST_REPORT"): + Path(output).with_suffix( + f".rank{torch.distributed.get_rank()}.json" + ).write_text(json.dumps(rows, indent=2)) + finally: + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + test_sft_packing_loss_and_gradients() diff --git a/tests/integration/megatron/test_shared_expert_stream_handoff.py b/tests/integration/megatron/test_shared_expert_stream_handoff.py new file mode 100644 index 000000000..a41b5bdd4 --- /dev/null +++ b/tests/integration/megatron/test_shared_expert_stream_handoff.py @@ -0,0 +1,44 @@ +"""Shared-expert output may not be consumed before its side stream finishes.""" + +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compiled_shared_expert_handoff(): + from megatron.core.transformer.moe.shared_experts import SharedExpertMLP + + from art.megatron.compile_workarounds import ( + _install_shared_expert_handoff_workaround, + ) + + _install_shared_expert_handoff_workaround() + + @torch.compiler.disable + def opaque_producer_delay(): + # Reproduce the graph break at a TE call without loading a model. + torch.cuda._sleep(10_000_000) + + class Handoff(torch.nn.Module): + get_output = SharedExpertMLP.get_output + + def __init__(self): + super().__init__() + self.config = SimpleNamespace(moe_shared_expert_overlap=True) + self.use_shared_expert_gate = False + self.stream = torch.cuda.Stream() + self.cached_output = None + + def forward(self, x): + self.stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.stream): + opaque_producer_delay() + self.cached_output = x + 1 + return self.get_output() * 2 + + model = torch.compile(Handoff()) + for value in range(1, 6): + inputs = torch.full((193, 2688), float(value), device="cuda") + torch.testing.assert_close(model(inputs), 2 * (inputs + 1), rtol=0, atol=0)