From 4f33faba6744601de1578625792763adf8d1b44b Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 20:02:28 +0000 Subject: [PATCH 1/3] fix(trainer-rank): expose source positions for context-parallel outputs --- src/art/trainer_rank/__init__.py | 12 ++ src/art/trainer_rank/_impl.py | 28 +++- .../unit/test_trainer_rank_head_recompute.py | 123 +++++++++++++++++- tests/unit/test_trainer_rank_moe_memory.py | 4 +- tests/unit/test_trainer_rank_split_peak.py | 8 +- tests/unit/test_trainer_rank_validation.py | 9 +- 6 files changed, 169 insertions(+), 15 deletions(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index 340d71560..eb8f588e4 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -253,6 +253,15 @@ def forward_micro_batches( ART moves its packed model inputs and labels internally without mutating the caller-owned `ForwardInput` objects. + With context parallelism, per-position outputs are local sequence shards, + not full sequences or necessarily equal contiguous splits. Each active + `ForwardOutput.positions` maps its rows to the request's flattened input + positions, including for hidden states, logits, top-k and target logprobs. + Index full-sequence masks/labels by these positions before using them. + For global means/pools, reduce local sums and counts with `dp_reduce` + (which includes context-parallel ranks), then divide; do not average + local means. No context-parallel output gather is performed. + Empty local microbatches are skipped unless `yield_empty=True`. Every rank must use the same setting. When a wave skips ranks, TrainerRank collective methods raise if called from its loop body; fully populated @@ -333,6 +342,9 @@ def dp_rank_forward( ) -> ForwardOutputs: """Forward inputs already local to this data-parallel rank. + Outputs remain context-parallel-local; use `ForwardOutput.positions` + to align masks and readouts as described in `forward_micro_batches`. + Per-input checkpoints and `no_grad` values override the method defaults. `no_grad=None` inherits the ambient PyTorch grad mode; `True` disables grads and `False` enables them. diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index de15bbdee..ace9f2c3e 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -184,12 +184,27 @@ class _LocalLoRASlotRef: @dataclass(frozen=True) class ForwardOutput(Generic[LogprobsT, TopKT, LogitsT, HiddenStatesT]): + """Per-request results, local to this context-parallel rank. + + ``positions`` maps axis 0 of every returned tensor (including both ``top_k`` + tensors) to offsets in the request's flattened ``input_tokens``. It is a + 1-D int64 tensor on the output device: ``arange(input_tokens.numel())`` + without context parallelism, and possibly empty/noncontiguous with it. + ``None`` denotes a request with no requested output fields. + + Select full-sequence masks/labels with ``index_select(0, output.positions)`` + before applying them to local outputs. Positions do not apply a next-token + shift; any alignment with target tokens is supplied by the caller. Exclude + the first source token with ``positions != 0``, not by slicing a local row. + """ + target_logprobs: LogprobsT top_k: TopKT logits: LogitsT hidden_states: HiddenStatesT checkpoint: str | None = None no_grad: bool = False + positions: torch.Tensor | None = None @dataclass(slots=True) @@ -4811,6 +4826,7 @@ def track(tensor: torch.Tensor | None) -> torch.Tensor | None: hidden_states=track(output.hidden_states), checkpoint=output.checkpoint, no_grad=output.no_grad, + positions=output.positions, ) for output in outputs ] @@ -4943,7 +4959,7 @@ def _estimate_group_request_output_bytes( self, requests: Sequence[AnyForwardInput], ) -> int: - total = 0 + total = _active_logical_tokens(requests) * _dtype_size(torch.long) for request in requests: seq_len = int(request.input_tokens.numel()) if request.target_tokens is not None: @@ -5349,6 +5365,10 @@ def _project_head( else None ) device = hidden_by_row.device + source_positions = tuple( + positions.to(device=device) + for positions in prepared.source_positions_by_item + ) target_logprobs = [None for _ in items] logits: list[torch.Tensor | None] = [None for _ in items] top_k: list[TopK | None] = [None for _ in items] @@ -5362,8 +5382,9 @@ def _project_head( if item.request.logits or item.request.top_k is not None: projected_rows.append(positions) if item.labels is not None: - source_positions = prepared.source_positions_by_item[index].to(device) - labels = item.labels.to(device=device).index_select(0, source_positions) + labels = item.labels.to(device=device).index_select( + 0, source_positions[index] + ) label_rows[index] = labels target_logprobs[index] = torch.zeros( tuple(labels.shape), @@ -5453,6 +5474,7 @@ def _project_head( if item.request.hidden_states else None ), + positions=source_positions[index], ) for index, (item, positions) in enumerate( zip(items, prepared.positions_by_item, strict=True) diff --git a/tests/unit/test_trainer_rank_head_recompute.py b/tests/unit/test_trainer_rank_head_recompute.py index bcf94df97..6cbabe4bb 100644 --- a/tests/unit/test_trainer_rank_head_recompute.py +++ b/tests/unit/test_trainer_rank_head_recompute.py @@ -7,6 +7,7 @@ import pytest import torch +from art.megatron.prefix_tree_packing import prefix_tree_pack from art.trainer_rank import ForwardInput, TrainerRank, _impl @@ -25,12 +26,8 @@ def _direct(function, *args, use_reentrant=False, **kwargs): return function(*args, **kwargs) -@pytest.mark.parametrize("top_k", (None, 3, 12)) -@pytest.mark.parametrize("include_logits", (False, True)) -@pytest.mark.parametrize("multi_target", (False, True)) -def test_recomputed_head_preserves_outputs_and_arbitrary_loss_gradients( - monkeypatch, top_k, include_logits, multi_target -): +@pytest.fixture +def local_head(monkeypatch): monkeypatch.setattr(_impl, "_HEAD_CHUNK_TOKENS", 16) monkeypatch.setattr(_impl, "_language_model", lambda model: model) monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda value: value) @@ -46,6 +43,14 @@ def test_recomputed_head_preserves_outputs_and_arbitrary_loss_gradients( monkeypatch.setattr( TrainerRank, "_gather_tensor_parallel_logits", lambda self, value: value ) + + +@pytest.mark.parametrize("top_k", (None, 3, 12)) +@pytest.mark.parametrize("include_logits", (False, True)) +@pytest.mark.parametrize("multi_target", (False, True)) +def test_recomputed_head_preserves_outputs_and_arbitrary_loss_gradients( + local_head, top_k, include_logits, multi_target +): generator = torch.Generator().manual_seed(11) weight = torch.randn(65, 7, generator=generator) / 3 hidden = torch.randn(41, 7, generator=generator) @@ -134,3 +139,109 @@ def pack(tensor): assert not any( shape[-1:] == (65,) and len(shape) == 2 for shape, _, _ in recomputed[4] ) + + +@pytest.mark.parametrize("cp_size", (1, 2, 4)) +@pytest.mark.parametrize("fields", ("hidden", "targets", "all")) +def test_output_positions_align_sharded_fields_masks_and_gradients( + local_head, cp_size, fields +): + tokens = [torch.tensor(row) for row in ([1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 8, 9])] + packed = prefix_tree_pack(tokens, max_depth=1) + generator = torch.Generator().manual_seed(911) + model = SimpleNamespace( + output_layer=_Head(torch.randn(11, 5, generator=generator)), + vocab_size=11, + share_embeddings_and_output_weights=False, + _scale_logits=lambda value: value, + ) + trainer = object.__new__(TrainerRank) + trainer.runtime = SimpleNamespace(model=[model]) + hidden = torch.randn(packed.tokens.numel(), 5, generator=generator).requires_grad_() + requests = [ + ForwardInput( + input_tokens=row, + target_tokens=( + torch.stack((row, row.roll(1)), dim=1) if fields != "hidden" else None + ), + hidden_states=fields != "targets", + logits=fields == "all", + top_k=3 if fields == "all" else None, + ) + for row in tokens + ] + items = [trainer._forward_item(request) for request in requests] + full = trainer._project_head( + items, + SimpleNamespace( + positions_by_item=packed.positions_by_sequence, + source_positions_by_item=tuple(torch.arange(len(row)) for row in tokens), + ), + hidden, + ) + for output, row in zip(full, tokens, strict=True): + torch.testing.assert_close(output.positions, torch.arange(len(row))) + + # Uneven, noncontiguous ownership; the last CP4 rank owns no source rows. + owners = torch.tensor([0, 1, 1, 2, 0, 1, 1, 2, 0]) % cp_size + seen = [[] for _ in items] + local_loss = hidden.sum() * 0 + full_loss = hidden.sum() * 0 + for output, row in zip(full, tokens, strict=True): + values = output.hidden_states if fields == "hidden" else output.target_logprobs + full_loss = full_loss + values[row % 2 == 1].square().sum() + for cp_rank in range(cp_size): + rows = torch.nonzero(owners == cp_rank).flatten().flip(0) + # Include a padding row that must never appear in public positions. + dispatched = torch.cat((rows, torch.tensor([-1]))) + pairs = [ + _impl._local_position_pairs(dispatched[None], positions) + for positions in packed.positions_by_sequence + ] + local_hidden = torch.cat((hidden[rows], hidden.new_zeros((1, 5)))) + outputs = trainer._project_head( + items, + SimpleNamespace( + positions_by_item=tuple(pair[0] for pair in pairs), + source_positions_by_item=tuple(pair[1] for pair in pairs), + ), + local_hidden, + ) + for index, (output, reference, row) in enumerate( + zip(outputs, full, tokens, strict=True) + ): + positions = output.positions + assert positions is not None and positions.dtype == torch.long + assert positions.device == local_hidden.device + assert not positions.requires_grad + torch.testing.assert_close( + row[positions], packed.tokens.flatten()[dispatched[pairs[index][0]]] + ) + seen[index].extend(positions.tolist()) + for field in ("hidden_states", "target_logprobs", "logits"): + actual, expected = getattr(output, field), getattr(reference, field) + if expected is not None: + torch.testing.assert_close(actual, expected[positions]) + if output.top_k is not None: + torch.testing.assert_close( + output.top_k.tokens, reference.top_k.tokens[positions] + ) + torch.testing.assert_close( + output.top_k.logprobs, reference.top_k.logprobs[positions] + ) + mask = (row % 2 == 1).index_select(0, positions) + values = ( + output.hidden_states if fields == "hidden" else output.target_logprobs + ) + local_loss = local_loss + values[mask].square().sum() + assert [sorted(positions) for positions in seen] == [ + list(range(len(row))) for row in tokens + ] + torch.testing.assert_close(local_loss, full_loss) + parameters = ( + (hidden, model.output_layer.weight) if fields != "hidden" else (hidden,) + ) + torch.testing.assert_close( + torch.autograd.grad(local_loss, parameters), + torch.autograd.grad(full_loss, parameters), + ) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 03f3b6d3c..308827de4 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -272,8 +272,8 @@ def test_summed_group_envelope_and_retained_profile_unchanged(): ) assert len(plan.groups) == 2 assert plan.packed_tokens == 16 - assert plan.output_bytes == 16 * 4 - assert rank._plan_cost(plan).required == int((16 * 65536 + 16 * 4) * 1.1) + assert plan.output_bytes == 16 * (4 + 8) + assert rank._plan_cost(plan).required == int((16 * 65536 + 16 * (4 + 8)) * 1.1) plan = replace(plan, packed_tokens=200, logical_tokens=200, output_bytes=4000) required = rank._plan_cost(plan).required assert required == int((200 * 65536 + 4000) * 1.1) diff --git a/tests/unit/test_trainer_rank_split_peak.py b/tests/unit/test_trainer_rank_split_peak.py index 7eb6cab04..065a92994 100644 --- a/tests/unit/test_trainer_rank_split_peak.py +++ b/tests/unit/test_trainer_rank_split_peak.py @@ -151,7 +151,9 @@ def test_profile_order_change_cannot_drop_completed_split_floor(monkeypatch): monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _: 10_100) rank._record_split_memory_floor(plan, 100, 1_200) # The real profile update changes only cost/order, not requests or geometry. - rank._update_memory_profile(b, 3_000, retained_bytes=500) + rank._update_memory_profile( + b, b.output_bytes + 2_600, retained_bytes=b.output_bytes + 100 + ) assert rank._plan_cost(b).ephemeral > rank._plan_cost(a).ephemeral before = dict(rank._memory_profiles) monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 10_000) @@ -177,7 +179,9 @@ def _counter_split(monkeypatch): 100, retained_compute_bytes_per_token=0, ) - monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 10_000) + monkeypatch.setattr( + rank, "_available_memory_bytes", lambda: 10_000 + 2 * child.output_bytes + ) counters: dict[str, Any] = dict(allocated=100, peak=100, resets=[], executed=0) monkeypatch.setattr(tr, "_telemetry_phase", lambda *a, **k: nullcontext()) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index d6ed71048..bd874b589 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3078,12 +3078,14 @@ def test_trainer_rank_graph_tracking_does_not_copy_outputs() -> None: trainer = TrainerRank(_runtime()) ref = _slot_ref("teacher") source = torch.ones(4, requires_grad=True) * 2 - output = ForwardOutput(source, None, None, None) + positions = torch.tensor([0, 2, 7, 8]) + output = ForwardOutput(source, None, None, None, positions=positions) tracked = trainer._track_slot_graph_outputs(ref, [output])[0] assert tracked.target_logprobs is not None assert tracked.target_logprobs.data_ptr() == source.data_ptr() + assert tracked.positions is positions def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None: @@ -3821,8 +3823,11 @@ def _preprocess(self, *args: object, **kwargs: object) -> None: topk_bytes = 3 * 5 * (4 + 8) logits_bytes = 3 * 10 * 4 hidden_bytes = 3 * 4 * 4 + positions_bytes = 3 * 8 assert estimate is not None and estimate[0] == plan.packed_tokens - assert plan.output_bytes == target_bytes + topk_bytes + logits_bytes + hidden_bytes + assert plan.output_bytes == ( + target_bytes + topk_bytes + logits_bytes + hidden_bytes + positions_bytes + ) def test_disconnected_outputs_keep_zero_graph_anchor() -> None: From 855736f3ef2a7af03a13edad5e7ac4d517b92356 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 20:26:16 +0000 Subject: [PATCH 2/3] fix(trainer-rank): reconstruct full context-parallel outputs internally --- dev/trainer_rank_check.py | 30 +- scripts/ci/trainer-rank-gpu-tests.sh | 1 + src/art/trainer_rank/__init__.py | 18 +- src/art/trainer_rank/_impl.py | 130 ++++++-- .../megatron/lora/test_dynamic_lora_slots.py | 2 +- .../unit/test_trainer_rank_head_recompute.py | 288 ++++++++++++------ tests/unit/test_trainer_rank_moe_memory.py | 32 +- tests/unit/test_trainer_rank_split_peak.py | 8 +- tests/unit/test_trainer_rank_validation.py | 9 +- 9 files changed, 341 insertions(+), 177 deletions(-) diff --git a/dev/trainer_rank_check.py b/dev/trainer_rank_check.py index d71784618..b0acf6c42 100644 --- a/dev/trainer_rank_check.py +++ b/dev/trainer_rank_check.py @@ -226,34 +226,10 @@ def _local_outputs( rank: TrainerRank, indexed_requests: Sequence[tuple[int, ForwardInput]], ) -> list[dict[str, object]]: - from art.megatron.lora import use_lora_slot - - requests = [request for _, request in indexed_requests] - plan = rank._plan_flat_forward(requests) - outputs: list[ForwardOutput] = [ - ForwardOutput(None, None, None, None) for _ in requests - ] - sources: list[torch.Tensor] = [torch.empty(0, dtype=torch.long) for _ in requests] - for group in plan.groups: - prepared = rank._prepare_packed_forward(group.packed) - with use_lora_slot(group.slot_ref): - group_outputs = rank._forward_packed(group.items, prepared) - for index, source, output in zip( - group.request_indices, - prepared.source_positions_by_item, - group_outputs, - strict=True, - ): - sources[index] = source - outputs[index] = output + outputs = rank.dp_rank_forward([request for _, request in indexed_requests]) return [ - _output_record(global_index, source, output) - for (global_index, _), source, output in zip( - indexed_requests, - sources, - outputs, - strict=True, - ) + _output_record(index, torch.arange(request.input_tokens.numel()), output) + for (index, request), output in zip(indexed_requests, outputs, strict=True) ] diff --git a/scripts/ci/trainer-rank-gpu-tests.sh b/scripts/ci/trainer-rank-gpu-tests.sh index b496c80b4..690d99db3 100755 --- a/scripts/ci/trainer-rank-gpu-tests.sh +++ b/scripts/ci/trainer-rank-gpu-tests.sh @@ -9,6 +9,7 @@ runtime_python="$( test -x "${runtime_python}" "${runtime_python}" -m pytest --tb=short \ + tests/unit/test_trainer_rank_head_recompute.py \ tests/unit/test_trainer_rank_custom_tensors.py \ tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py \ 'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_packed_sibling_order_matches_cp1_oracle[2]' \ diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index eb8f588e4..4d666d1ff 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -253,14 +253,11 @@ def forward_micro_batches( ART moves its packed model inputs and labels internally without mutating the caller-owned `ForwardInput` objects. - With context parallelism, per-position outputs are local sequence shards, - not full sequences or necessarily equal contiguous splits. Each active - `ForwardOutput.positions` maps its rows to the request's flattened input - positions, including for hidden states, logits, top-k and target logprobs. - Index full-sequence masks/labels by these positions before using them. - For global means/pools, reduce local sums and counts with `dp_reduce` - (which includes context-parallel ranks), then divide; do not average - local means. No context-parallel output gather is performed. + Per-position outputs contain the full flattened input sequence in source + order, including with context parallelism. TP/CP ranks compute the same + loss on these replicated outputs; ART routes gradients to owning rows + without multiplying them by the number of replicas. `dp_reduce` combines + only distinct data-parallel batches. Empty local microbatches are skipped unless `yield_empty=True`. Every rank must use the same setting. When a wave skips ranks, TrainerRank @@ -342,8 +339,8 @@ def dp_rank_forward( ) -> ForwardOutputs: """Forward inputs already local to this data-parallel rank. - Outputs remain context-parallel-local; use `ForwardOutput.positions` - to align masks and readouts as described in `forward_micro_batches`. + Outputs contain full sequences in source order on every TP/CP rank, + with the same loss and reduction contract as `forward_micro_batches`. Per-input checkpoints and `no_grad` values override the method defaults. `no_grad=None` inherits the ambient PyTorch grad mode; `True` disables @@ -364,6 +361,7 @@ def dp_reduce( *, op: dist.ReduceOp.RedOpType = dist.ReduceOp.SUM, ) -> None: + """Reduce in place over data-parallel batches, excluding TP/CP replicas.""" super().dp_reduce(tensor, op=op) def optim_step( diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index ace9f2c3e..c0ff00e38 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -184,19 +184,7 @@ class _LocalLoRASlotRef: @dataclass(frozen=True) class ForwardOutput(Generic[LogprobsT, TopKT, LogitsT, HiddenStatesT]): - """Per-request results, local to this context-parallel rank. - - ``positions`` maps axis 0 of every returned tensor (including both ``top_k`` - tensors) to offsets in the request's flattened ``input_tokens``. It is a - 1-D int64 tensor on the output device: ``arange(input_tokens.numel())`` - without context parallelism, and possibly empty/noncontiguous with it. - ``None`` denotes a request with no requested output fields. - - Select full-sequence masks/labels with ``index_select(0, output.positions)`` - before applying them to local outputs. Positions do not apply a next-token - shift; any alignment with target tokens is supplied by the caller. Exclude - the first source token with ``positions != 0``, not by slicing a local row. - """ + """Per-request tensors in flattened input order, replicated across TP/CP.""" target_logprobs: LogprobsT top_k: TopKT @@ -204,7 +192,6 @@ class ForwardOutput(Generic[LogprobsT, TopKT, LogitsT, HiddenStatesT]): hidden_states: HiddenStatesT checkpoint: str | None = None no_grad: bool = False - positions: torch.Tensor | None = None @dataclass(slots=True) @@ -629,6 +616,34 @@ def backward( return grad_outputs[0], None +class _GatherContextParallelRows(torch.autograd.Function): + @staticmethod + def forward( + ctx: FunctionCtx, + tensor: torch.Tensor, + positions: torch.Tensor, + length: int, + group: dist.ProcessGroup, + ) -> torch.Tensor: + ctx.save_for_backward(positions) + # Each source row has one owner. Scatter + SUM handles unequal/empty + # shards without padded all-gather buffers, including for full logits. + output = tensor.new_zeros((length, *tensor.shape[1:])) + output.index_copy_(0, positions, tensor) + dist.all_reduce(output, group=group) + return output + + @staticmethod + def backward( + ctx: FunctionCtx, *grad_outputs: torch.Tensor + ) -> tuple[torch.Tensor, None, None, None]: + (positions,) = cast(tuple[torch.Tensor, ...], getattr(ctx, "saved_tensors")) + # The caller's loss is replicated over CP, just as over TP after the + # sequence-parallel gather with tensor_parallel_output_grad=False. + # Route one copy to its owner, without summing CP copies of the loss. + return grad_outputs[0].index_select(0, positions), None, None, None + + class _CustomSlotGraphSentinel(torch.autograd.Function): @staticmethod def forward( @@ -881,6 +896,7 @@ class _PreparedPackedForward: packed_seq_params: "PackedSeqParams | None" positions_by_item: tuple[torch.Tensor, ...] source_positions_by_item: tuple[torch.Tensor, ...] + context_parallel_group: dist.ProcessGroup | None = None type _RowMatch = tuple[torch.Tensor, torch.Tensor, tuple[int, ...]] @@ -3027,7 +3043,7 @@ def dp_reduce( dist.all_reduce( tensor, op=op, - group=ps.get_data_parallel_group(with_context_parallel=True), + group=ps.get_data_parallel_group(with_context_parallel=False), ) def optim_step( @@ -3799,6 +3815,10 @@ def add( for param, grad in zip(params, grads, strict=True): if bool(getattr(param, "allreduce", True)): group = ps.get_data_parallel_group(with_context_parallel=True) + if getattr(param, "_art_custom_checkpoint_param", False): + # Custom heads consume replicated full-sequence outputs; + # average their CP copies while still summing DP batches. + grad.div_(ps.get_context_parallel_world_size()) else: group = ps.get_expert_data_parallel_group() if group is not None and group.size() > 1: @@ -4826,7 +4846,6 @@ def track(tensor: torch.Tensor | None) -> torch.Tensor | None: hidden_states=track(output.hidden_states), checkpoint=output.checkpoint, no_grad=output.no_grad, - positions=output.positions, ) for output in outputs ] @@ -4959,7 +4978,7 @@ def _estimate_group_request_output_bytes( self, requests: Sequence[AnyForwardInput], ) -> int: - total = _active_logical_tokens(requests) * _dtype_size(torch.long) + total = 0 for request in requests: seq_len = int(request.input_tokens.numel()) if request.target_tokens is not None: @@ -5111,6 +5130,10 @@ def _estimate_required_memory_bytes_from_values( static_compute = max( static_compute, packed_tokens * self._moe_output_bytes_per_token ) + if signature.topology[2] > 1: + # Local head results coexist with full CP outputs during gathering. + # Uneven rank plans can assign all of an item's rows to one rank. + static_compute += output_bytes # A profile learned under lighter sharing (lower logical/packed ratio) # underestimates the per-packed-token footprint of a deeper-shared # plan; scale the trusted estimate up by the ratio gap. @@ -5303,7 +5326,67 @@ def _forward_packed( hidden_by_row = self._gather_sequence_parallel_hidden( self._decoder_hidden(prepared) ) - return self._project_head(items, prepared, hidden_by_row) + outputs = self._project_head(items, prepared, hidden_by_row) + group = prepared.context_parallel_group + if group is None: + return outputs + tensors = [ + tensor + for output in outputs + for tensor in ( + output.target_logprobs, + output.logits, + output.hidden_states, + None if output.top_k is None else output.top_k.logprobs, + None if output.top_k is None else output.top_k.tokens, + ) + if tensor is not None + ] + grad_flags = [False for _ in tensors] + if torch.is_grad_enabled(): + flags = torch.tensor( + [ + tensor.is_floating_point() + and (tensor.requires_grad or hidden_by_row.requires_grad) + for tensor in tensors + ], + device=hidden_by_row.device, + dtype=torch.int32, + ) + dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=group) + grad_flags = flags.tolist() + needs_grad = iter(grad_flags) + + def gather(tensor: torch.Tensor | None) -> torch.Tensor | None: + if tensor is None: + return None + if next(needs_grad) and not tensor.requires_grad: + # Empty shards must still enter decoder backward collectives. + # A frozen decoder with a trainable head only needs a leaf. + tensor = tensor + hidden_by_row.reshape(-1)[:1].sum() * 0.0 + tensor.requires_grad_(True) + return _GatherContextParallelRows.apply(tensor, positions, length, group) + + for index, (item, output, source_positions) in enumerate( + zip(items, outputs, prepared.source_positions_by_item, strict=True) + ): + positions = source_positions.to(device=hidden_by_row.device) + length = int(item.input_ids.numel()) + outputs[index] = replace( + output, + target_logprobs=gather(output.target_logprobs), + logits=gather(output.logits), + hidden_states=gather(output.hidden_states), + top_k=( + TopK( + cast(torch.Tensor, gather(output.top_k.logprobs)), + cast(torch.Tensor, gather(output.top_k.tokens)), + ) + if output.top_k is not None + else None + ), + ) + return outputs def _decoder_hidden( self, @@ -5365,10 +5448,6 @@ def _project_head( else None ) device = hidden_by_row.device - source_positions = tuple( - positions.to(device=device) - for positions in prepared.source_positions_by_item - ) target_logprobs = [None for _ in items] logits: list[torch.Tensor | None] = [None for _ in items] top_k: list[TopK | None] = [None for _ in items] @@ -5382,9 +5461,8 @@ def _project_head( if item.request.logits or item.request.top_k is not None: projected_rows.append(positions) if item.labels is not None: - labels = item.labels.to(device=device).index_select( - 0, source_positions[index] - ) + source_positions = prepared.source_positions_by_item[index].to(device) + labels = item.labels.to(device=device).index_select(0, source_positions) label_rows[index] = labels target_logprobs[index] = torch.zeros( tuple(labels.shape), @@ -5474,7 +5552,6 @@ def _project_head( if item.request.hidden_states else None ), - positions=source_positions[index], ) for index, (item, positions) in enumerate( zip(items, prepared.positions_by_item, strict=True) @@ -5921,6 +5998,7 @@ def _prepare_context_parallel_forward( packed_seq_params=prepared.packed_seq_params, positions_by_item=tuple(pair[0] for pair in local_position_pairs), source_positions_by_item=tuple(pair[1] for pair in local_position_pairs), + context_parallel_group=ps.get_context_parallel_group(), ) def _topology(self) -> "ParallelTopology": diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py index 140e1fef3..cd9d47f60 100644 --- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py +++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py @@ -272,7 +272,7 @@ def _custom_parameter_reduction_worker( torch.testing.assert_close(parameter, torch.tensor(1.0, device=device)) (parameter * float(rank + 1)).backward() (reduced,) = trainer._reduce_dynamic_grads((parameter,), scale_grads=1.0) - expected = {"dp": 3.0, "tp": 1.5, "cp": 3.0, "tp_cp": 5.0}[topology] + expected = {"dp": 3.0, "tp": 1.5, "cp": 1.5, "tp_cp": 2.5}[topology] torch.testing.assert_close(reduced, torch.tensor(expected, device=device)) finally: if getattr(ps, "model_parallel_is_initialized", lambda: False)(): diff --git a/tests/unit/test_trainer_rank_head_recompute.py b/tests/unit/test_trainer_rank_head_recompute.py index 6cbabe4bb..da348a64c 100644 --- a/tests/unit/test_trainer_rank_head_recompute.py +++ b/tests/unit/test_trainer_rank_head_recompute.py @@ -1,11 +1,14 @@ from __future__ import annotations from contextlib import nullcontext +from datetime import timedelta from types import SimpleNamespace from unittest.mock import patch import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp from art.megatron.prefix_tree_packing import prefix_tree_pack from art.trainer_rank import ForwardInput, TrainerRank, _impl @@ -26,8 +29,7 @@ def _direct(function, *args, use_reentrant=False, **kwargs): return function(*args, **kwargs) -@pytest.fixture -def local_head(monkeypatch): +def _patch_local_head(monkeypatch): monkeypatch.setattr(_impl, "_HEAD_CHUNK_TOKENS", 16) monkeypatch.setattr(_impl, "_language_model", lambda model: model) monkeypatch.setattr(_impl, "_all_reduce_tensor_parallel_max", lambda value: value) @@ -49,8 +51,9 @@ def local_head(monkeypatch): @pytest.mark.parametrize("include_logits", (False, True)) @pytest.mark.parametrize("multi_target", (False, True)) def test_recomputed_head_preserves_outputs_and_arbitrary_loss_gradients( - local_head, top_k, include_logits, multi_target + monkeypatch, top_k, include_logits, multi_target ): + _patch_local_head(monkeypatch) generator = torch.Generator().manual_seed(11) weight = torch.randn(65, 7, generator=generator) / 3 hidden = torch.randn(41, 7, generator=generator) @@ -141,107 +144,196 @@ def pack(tensor): ) -@pytest.mark.parametrize("cp_size", (1, 2, 4)) -@pytest.mark.parametrize("fields", ("hidden", "targets", "all")) -def test_output_positions_align_sharded_fields_masks_and_gradients( - local_head, cp_size, fields +@pytest.mark.parametrize("cp_size,dp_size", ((2, 1), (4, 1), (2, 2))) +def test_context_parallel_outputs_match_full_sequence(cp_size, dp_size, tmp_path): + mp.spawn( + _context_parallel_worker, + args=(cp_size, dp_size, f"file://{tmp_path / 'cp'}", "gloo"), + nprocs=cp_size * dp_size, + join=True, + ) + + +@pytest.mark.parametrize("cp_size", (2, 4)) +def test_context_parallel_outputs_cuda(cp_size, tmp_path): + if not torch.cuda.is_available() or torch.cuda.device_count() < cp_size: + pytest.skip(f"requires {cp_size} CUDA devices") + mp.spawn( + _context_parallel_worker, + args=(cp_size, 1, f"file://{tmp_path / 'cp'}", "nccl"), + nprocs=cp_size, + join=True, + ) + + +def _context_parallel_worker(rank, cp_size, dp_size, init_method, backend): + from megatron.core import parallel_state as ps + + device = torch.device("cpu" if backend == "gloo" else f"cuda:{rank}") + if device.type == "cuda": + torch.cuda.set_device(device) + dist.init_process_group( + backend, + init_method=init_method, + rank=rank, + world_size=cp_size * dp_size, + timeout=timedelta(seconds=90), + ) + try: + cp_groups = [ + dist.new_group(list(range(dp * cp_size, (dp + 1) * cp_size))) + for dp in range(dp_size) + ] + dp_groups = [ + dist.new_group(list(range(cp, cp_size * dp_size, cp_size))) + for cp in range(cp_size) + ] + cp_rank, dp_rank = rank % cp_size, rank // cp_size + cp_group, dp_group = cp_groups[dp_rank], dp_groups[cp_rank] + with pytest.MonkeyPatch.context() as monkeypatch: + _patch_local_head(monkeypatch) + monkeypatch.setattr(ps, "get_context_parallel_world_size", lambda: cp_size) + monkeypatch.setattr( + ps, + "get_data_parallel_group", + lambda *, with_context_parallel: ( + dist.group.WORLD if with_context_parallel else dp_group + ), + ) + monkeypatch.setattr(ps, "get_tensor_model_parallel_group", lambda **_: None) + for mode in ( + "hidden", + "targets", + "all", + "logits", + "head_only", + "frozen", + "no_grad", + ): + _check_context_parallel_case( + cp_rank, dp_rank, cp_size, dp_size, cp_group, device, mode + ) + finally: + dist.destroy_process_group() + + +def _check_context_parallel_case( + cp_rank, dp_rank, cp_size, dp_size, cp_group, device, mode ): tokens = [torch.tensor(row) for row in ([1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 8, 9])] packed = prefix_tree_pack(tokens, max_depth=1) - generator = torch.Generator().manual_seed(911) - model = SimpleNamespace( - output_layer=_Head(torch.randn(11, 5, generator=generator)), - vocab_size=11, - share_embeddings_and_output_weights=False, - _scale_logits=lambda value: value, - ) - trainer = object.__new__(TrainerRank) - trainer.runtime = SimpleNamespace(model=[model]) - hidden = torch.randn(packed.tokens.numel(), 5, generator=generator).requires_grad_() - requests = [ - ForwardInput( - input_tokens=row, - target_tokens=( - torch.stack((row, row.roll(1)), dim=1) if fields != "hidden" else None - ), - hidden_states=fields != "targets", - logits=fields == "all", - top_k=3 if fields == "all" else None, + generator = torch.Generator(device=device).manual_seed(911) + features = torch.randn(9, 3, generator=generator, device=device) + dp_rank / 5 + decoder_weight = torch.randn(3, 5, generator=generator, device=device) + head_weight = torch.randn(11, 5, generator=generator, device=device) + probe_weight = torch.randn(5, 2, generator=generator, device=device) + requests = [] + for index, row in enumerate(tokens): + labels = torch.stack((row, row.roll(1)), dim=1) + labels[::2] = -100 + if index == 1: + labels.fill_(-100) + requests.append( + ForwardInput( + input_tokens=row, + target_tokens=labels if mode not in ("hidden", "logits") else None, + hidden_states=mode not in ("targets", "logits"), + logits=mode not in ("hidden", "targets"), + top_k=3 if mode not in ("hidden", "targets", "logits") else None, + ) ) - for row in tokens + # Unequal, reversed ownership with a shared prefix; CP4 rank 3 is empty. + owners = torch.tensor([0, 1, 1, 2, 0, 1, 1, 2, 0]) % cp_size + rows = torch.nonzero(owners == cp_rank).flatten().flip(0) + dispatched = torch.cat((rows, torch.tensor([-1]))) + pairs = [ + _impl._local_position_pairs(dispatched[None], positions) + for positions in packed.positions_by_sequence ] - items = [trainer._forward_item(request) for request in requests] - full = trainer._project_head( - items, - SimpleNamespace( - positions_by_item=packed.positions_by_sequence, - source_positions_by_item=tuple(torch.arange(len(row)) for row in tokens), - ), - hidden, - ) - for output, row in zip(full, tokens, strict=True): - torch.testing.assert_close(output.positions, torch.arange(len(row))) - # Uneven, noncontiguous ownership; the last CP4 rank owns no source rows. - owners = torch.tensor([0, 1, 1, 2, 0, 1, 1, 2, 0]) % cp_size - seen = [[] for _ in items] - local_loss = hidden.sum() * 0 - full_loss = hidden.sum() * 0 - for output, row in zip(full, tokens, strict=True): - values = output.hidden_states if fields == "hidden" else output.target_logprobs - full_loss = full_loss + values[row % 2 == 1].square().sum() - for cp_rank in range(cp_size): - rows = torch.nonzero(owners == cp_rank).flatten().flip(0) - # Include a padding row that must never appear in public positions. - dispatched = torch.cat((rows, torch.tensor([-1]))) - pairs = [ - _impl._local_position_pairs(dispatched[None], positions) - for positions in packed.positions_by_sequence - ] - local_hidden = torch.cat((hidden[rows], hidden.new_zeros((1, 5)))) - outputs = trainer._project_head( - items, - SimpleNamespace( - positions_by_item=tuple(pair[0] for pair in pairs), - source_positions_by_item=tuple(pair[1] for pair in pairs), - ), - local_hidden, + def run(local): + decoder = torch.nn.Parameter( + decoder_weight.clone(), requires_grad=mode not in ("frozen", "head_only") ) - for index, (output, reference, row) in enumerate( - zip(outputs, full, tokens, strict=True) - ): - positions = output.positions - assert positions is not None and positions.dtype == torch.long - assert positions.device == local_hidden.device - assert not positions.requires_grad - torch.testing.assert_close( - row[positions], packed.tokens.flatten()[dispatched[pairs[index][0]]] - ) - seen[index].extend(positions.tolist()) - for field in ("hidden_states", "target_logprobs", "logits"): - actual, expected = getattr(output, field), getattr(reference, field) - if expected is not None: - torch.testing.assert_close(actual, expected[positions]) - if output.top_k is not None: - torch.testing.assert_close( - output.top_k.tokens, reference.top_k.tokens[positions] - ) - torch.testing.assert_close( - output.top_k.logprobs, reference.top_k.logprobs[positions] + head = _Head(head_weight) + head.weight.requires_grad_(mode != "frozen") + probe = torch.nn.Parameter(probe_weight.clone(), requires_grad=mode != "frozen") + trainer = object.__new__(TrainerRank) + trainer._skipped_forward_waves = {} + trainer.runtime = SimpleNamespace( + model=[ + SimpleNamespace( + output_layer=head, + vocab_size=11, + share_embeddings_and_output_weights=False, + _scale_logits=lambda value: value, ) - mask = (row % 2 == 1).index_select(0, positions) - values = ( - output.hidden_states if fields == "hidden" else output.target_logprobs + ] + ) + trainer._tag_custom_parameters((probe,)) + with torch.set_grad_enabled(mode != "no_grad"): + hidden = features @ decoder + if local: + hidden = torch.cat((hidden[rows], hidden.new_zeros((1, 5)))) + trainer._decoder_hidden = lambda _: hidden + trainer._gather_sequence_parallel_hidden = lambda value: value + outputs = trainer._forward_packed( + [trainer._forward_item(request) for request in requests], + SimpleNamespace( + positions_by_item=( + tuple(pair[0] for pair in pairs) + if local + else packed.positions_by_sequence + ), + source_positions_by_item=( + tuple(pair[1] for pair in pairs) + if local + else tuple(torch.arange(len(row)) for row in tokens) + ), + context_parallel_group=cp_group if local else None, + ), ) - local_loss = local_loss + values[mask].square().sum() - assert [sorted(positions) for positions in seen] == [ - list(range(len(row))) for row in tokens - ] - torch.testing.assert_close(local_loss, full_loss) - parameters = ( - (hidden, model.output_layer.weight) if fields != "hidden" else (hidden,) - ) - torch.testing.assert_close( - torch.autograd.grad(local_loss, parameters), - torch.autograd.grad(full_loss, parameters), - ) + tensors, terms = [], [] + for output, row in zip(outputs, tokens, strict=True): + assert not hasattr(output, "positions") + values = [output.target_logprobs, output.logits, output.hidden_states] + if output.top_k is not None: + values.append(output.top_k.logprobs) + tensors.append(output.top_k.tokens) + for value in values: + if value is not None: + assert value.shape[0] == len(row) + tensors.append(value) + terms.append(value.mean().square() + value.square().mean()) + if output.hidden_states is not None: + # The original failing caller expression needs no CP mapping. + selected = output.hidden_states[1:][(row[1:] % 2 == 1).to(device)] + terms.append((selected @ probe).square().mean()) + loss = torch.stack(terms).sum() + if mode not in ("frozen", "no_grad"): + loss.backward() + return trainer, tensors, loss.detach(), (decoder, head.weight, probe) + + reference = run(False) + actual = run(True) + for value, expected in zip(actual[1], reference[1], strict=True): + torch.testing.assert_close(value, expected) + assert value.requires_grad == expected.requires_grad + torch.testing.assert_close(actual[2], reference[2]) + expected_loss = reference[2].clone() + dist.all_reduce(expected_loss) + expected_loss /= cp_size + trainer = actual[0] + trainer.dp_reduce(actual[2]) + torch.testing.assert_close(actual[2], expected_loss) + count = torch.tensor(sum(len(row) for row in tokens), device=device) + trainer.dp_reduce(count) + assert count.item() == dp_size * sum(len(row) for row in tokens) + if mode in ("frozen", "no_grad"): + return + reduced = trainer._reduce_dynamic_grads(actual[3], scale_grads=0.5) + for grad, param in zip(reduced, reference[3], strict=True): + expected = torch.zeros_like(param) if param.grad is None else param.grad.clone() + dist.all_reduce(expected) + expected *= 0.5 / cp_size + torch.testing.assert_close(grad, expected, atol=2e-5, rtol=2e-5) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 308827de4..e0dabad67 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -86,6 +86,34 @@ def _signature(): return _MemorySignature((1, 1, 1, 1), (1, None), 1, (), True, (True,)) +def test_cp_memory_charges_local_and_gathered_outputs(): + rank = _rank() + signature = replace(_signature(), topology=(1, 1, 4, 1)) + output_bytes = 1 << 30 + estimate = rank._estimate_required_memory_bytes_from_values( + packed_tokens=16, output_bytes=output_bytes, signature=signature + ) + compute = 16 * 2048 * 2 * 14 + assert estimate == int((compute + 2 * output_bytes) * 1.1) + # A warm profile includes gather workspace already; do not add it twice. + rank._update_memory_profile( + SimpleNamespace( + signature=signature, + packed_tokens=16, + output_bytes=output_bytes, + active_logical_tokens=16, + ), + compute + 2 * output_bytes, + retained_bytes=None, + ) + assert ( + rank._estimate_required_memory_bytes_from_values( + packed_tokens=16, output_bytes=output_bytes, signature=signature + ) + == estimate + ) + + def test_supported_constructor_and_original_shape(layer): rank = _rank(layer) assert rank._moe_output_bytes_per_token == (512 + 3 * 2048) * 8 * 2 @@ -272,8 +300,8 @@ def test_summed_group_envelope_and_retained_profile_unchanged(): ) assert len(plan.groups) == 2 assert plan.packed_tokens == 16 - assert plan.output_bytes == 16 * (4 + 8) - assert rank._plan_cost(plan).required == int((16 * 65536 + 16 * (4 + 8)) * 1.1) + assert plan.output_bytes == 16 * 4 + assert rank._plan_cost(plan).required == int((16 * 65536 + 16 * 4) * 1.1) plan = replace(plan, packed_tokens=200, logical_tokens=200, output_bytes=4000) required = rank._plan_cost(plan).required assert required == int((200 * 65536 + 4000) * 1.1) diff --git a/tests/unit/test_trainer_rank_split_peak.py b/tests/unit/test_trainer_rank_split_peak.py index 065a92994..7eb6cab04 100644 --- a/tests/unit/test_trainer_rank_split_peak.py +++ b/tests/unit/test_trainer_rank_split_peak.py @@ -151,9 +151,7 @@ def test_profile_order_change_cannot_drop_completed_split_floor(monkeypatch): monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda _: 10_100) rank._record_split_memory_floor(plan, 100, 1_200) # The real profile update changes only cost/order, not requests or geometry. - rank._update_memory_profile( - b, b.output_bytes + 2_600, retained_bytes=b.output_bytes + 100 - ) + rank._update_memory_profile(b, 3_000, retained_bytes=500) assert rank._plan_cost(b).ephemeral > rank._plan_cost(a).ephemeral before = dict(rank._memory_profiles) monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 10_000) @@ -179,9 +177,7 @@ def _counter_split(monkeypatch): 100, retained_compute_bytes_per_token=0, ) - monkeypatch.setattr( - rank, "_available_memory_bytes", lambda: 10_000 + 2 * child.output_bytes - ) + monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 10_000) counters: dict[str, Any] = dict(allocated=100, peak=100, resets=[], executed=0) monkeypatch.setattr(tr, "_telemetry_phase", lambda *a, **k: nullcontext()) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index bd874b589..d6ed71048 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3078,14 +3078,12 @@ def test_trainer_rank_graph_tracking_does_not_copy_outputs() -> None: trainer = TrainerRank(_runtime()) ref = _slot_ref("teacher") source = torch.ones(4, requires_grad=True) * 2 - positions = torch.tensor([0, 2, 7, 8]) - output = ForwardOutput(source, None, None, None, positions=positions) + output = ForwardOutput(source, None, None, None) tracked = trainer._track_slot_graph_outputs(ref, [output])[0] assert tracked.target_logprobs is not None assert tracked.target_logprobs.data_ptr() == source.data_ptr() - assert tracked.positions is positions def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None: @@ -3823,11 +3821,8 @@ def _preprocess(self, *args: object, **kwargs: object) -> None: topk_bytes = 3 * 5 * (4 + 8) logits_bytes = 3 * 10 * 4 hidden_bytes = 3 * 4 * 4 - positions_bytes = 3 * 8 assert estimate is not None and estimate[0] == plan.packed_tokens - assert plan.output_bytes == ( - target_bytes + topk_bytes + logits_bytes + hidden_bytes + positions_bytes - ) + assert plan.output_bytes == target_bytes + topk_bytes + logits_bytes + hidden_bytes def test_disconnected_outputs_keep_zero_graph_anchor() -> None: From 1e068eed529cfb2cd18dcf8c0680c8c4b89d4297 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 20:39:31 +0000 Subject: [PATCH 3/3] test(trainer-rank): guard context-parallel tests without Megatron --- src/art/trainer_rank/__init__.py | 4 ++-- src/art/trainer_rank/_impl.py | 1 + tests/unit/test_trainer_rank_head_recompute.py | 2 ++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index 4d666d1ff..fd321a820 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -254,8 +254,8 @@ def forward_micro_batches( the caller-owned `ForwardInput` objects. Per-position outputs contain the full flattened input sequence in source - order, including with context parallelism. TP/CP ranks compute the same - loss on these replicated outputs; ART routes gradients to owning rows + order, including with context parallelism. Callers must compute identical + losses on every TP/CP replica; ART routes gradients to owning rows without multiplying them by the number of replicas. `dp_reduce` combines only distinct data-parallel batches. diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index c0ff00e38..0e0045d61 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -3040,6 +3040,7 @@ def dp_reduce( self._guard_forward_collective("dp_reduce") from megatron.core import parallel_state as ps + # Public outputs are CP-replicated; internal shard reductions still include CP. dist.all_reduce( tensor, op=op, diff --git a/tests/unit/test_trainer_rank_head_recompute.py b/tests/unit/test_trainer_rank_head_recompute.py index da348a64c..1a2c8d9a8 100644 --- a/tests/unit/test_trainer_rank_head_recompute.py +++ b/tests/unit/test_trainer_rank_head_recompute.py @@ -146,6 +146,7 @@ def pack(tensor): @pytest.mark.parametrize("cp_size,dp_size", ((2, 1), (4, 1), (2, 2))) def test_context_parallel_outputs_match_full_sequence(cp_size, dp_size, tmp_path): + pytest.importorskip("megatron.core") mp.spawn( _context_parallel_worker, args=(cp_size, dp_size, f"file://{tmp_path / 'cp'}", "gloo"), @@ -158,6 +159,7 @@ def test_context_parallel_outputs_match_full_sequence(cp_size, dp_size, tmp_path def test_context_parallel_outputs_cuda(cp_size, tmp_path): if not torch.cuda.is_available() or torch.cuda.device_count() < cp_size: pytest.skip(f"requires {cp_size} CUDA devices") + pytest.importorskip("megatron.core") mp.spawn( _context_parallel_worker, args=(cp_size, 1, f"file://{tmp_path / 'cp'}", "nccl"),