Repository navigation
Conversation
prepare_autosp_inputs() accepted seq_dim but only used it for mark_dynamic: the graph passes always read the sequence length from shape[1] and sliced dimension 1. Sequence-first [seq, batch] inputs therefore sharded the batch dimension, or failed because that dimension is static. Dynamo copies only the `tag` attribute into node metadata, so record seq_dim in the tag as (key, seq_dim). Tag lookup accepts both forms, and a bare key keeps the previous default of dimension 1. The sequence length and the slice now use the recorded dimension, and label_id is bounds-checked like the other prepared inputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
AutoSP sharded the labels before the model's causal shift. Each rank therefore computed a local mean loss, which caused three problems: - The last target of every shard was paired with a wrong token. - Shards were weighted equally even when their numbers of valid (non-ignored) tokens differed. - A shard with no valid tokens produced NaN. Labels are now sharded where the cross entropy target still has the full label shape, so the shift happens on the full sequence. The cross entropy is rewritten to reduction="sum". autosp::aggregate_loss then sums the loss and the valid-token count over the SP group and returns the full-sequence loss. It supports both "mean" and "sum" reductions, so the num_items_in_batch path of HF losses is covered too. The engine averages gradients over every rank, SP ranks included. Backward therefore scales the local gradient by sp_size / total_valid, which makes the averaged gradient equal to the full-sequence gradient. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
… calls Summing a half-precision cross entropy over a long shard overflows even when the requested mean does not, so the per-token losses are now accumulated in FP32 and the result is cast back to the logits dtype. Cross entropy calls that pass input/target as keywords are now recognized instead of falling back to sharding the labels before the shift. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
# Conflicts: # deepspeed/compile/custom_ops/all_to_all.py # deepspeed/compile/util.py # docs/code-docs/source/training.rst # tests/unit/v1/compile/test_compile_autosp.py
…he caller Scope this PR to loss and gradient aggregation; label alignment under an in-model causal shift moves to a follow-up PR. - groups: with a (data_parallel, sequence_parallel) mesh and no mpu, the sequence-data-parallel group is the whole world, so gradients are summed over SP ranks. Previously it fell back to the DP-only mesh dimension while the engine still divided by dp/sp. - initialize: create that mesh when the autosp pass is enabled with only sequence_parallel_size, so the implicit and explicit paths share the Ulysses convention (sum over SP, average over DP). - autosp::aggregate_loss no longer scales the gradient by sp_size. - engine: reject an mpu/mesh whose SP size differs from AutoSP's. - sp_compile: drop the shape-based search for the full-length labels; labels are sharded as given and must be shifted on the full sequence. - Tests: explicit mesh, an all-ignored shard, the SP-size mismatch, and the attention-mask gradient tests updated to the summed convention. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6f7c57b9a6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| logger.info(f"config to Initialize mesh device: {config}") | ||
| mesh_device = dist.initialize_mesh_device((config["data_parallel_size"], config["sequence_parallel_size"]), \ | ||
| ("data_parallel", "sequence_parallel")) | ||
| elif mpu is None and _autosp_sequence_parallel_size(config) > 1: |
There was a problem hiding this comment.
When passes contains autosp, sequence_parallel_size > 1, and ZeRO stage is 2 or 3, this branch installs an SP mesh during initialization, but DeepSpeedEngine.get_autosp_backend() later returns None and uses ordinary torch.compile. The model is therefore not sequence-sharded while the engine/ZeRO gradient reduction still sums over the world group and divides only by the DP size, scaling the normal data-parallel update by sequence_parallel_size; reject these stages before installing the mesh or leave normal DP groups in place for the fallback.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in the base PR #8457 (e17b176): deepspeed.initialize now raises when AutoSP (sequence_parallel_size > 1 with the autosp pass) is combined with ZeRO stage 2/3, before the SP mesh is installed, so it can no longer fall back to plain torch.compile with SP-scaled gradients. Covered by test_rejects_zero_stages_without_autosp[2,3]. This PR is rebased on it.
| valid_mask = gm.graph.call_function(operator.ne, args=(target_node, ignore_index)) | ||
| with gm.graph.inserting_after(valid_mask): | ||
| valid_tokens = gm.graph.call_function(torch.sum, args=(valid_mask, )) |
There was a problem hiding this comment.
When labels are floating class-probability targets (target.shape == input.shape), which F.cross_entropy supports, every class probability is unequal to the default ignore index, so this mask reports N*C valid tokens. Cross entropy's mean reduction divides by N/spatial positions, not by C as well, so aggregate_loss scales both the full loss and gradients down by the vocabulary size; detect this target form and count examples, or reject it, rather than counting target elements.
Useful? React with 👍 / 👎.
| global_loss = gm.graph.call_function(operator.getitem, args=(aggregated, 0)) | ||
| with gm.graph.inserting_after(global_loss): | ||
| # Return the loss in the dtype cross entropy would have produced. | ||
| output_loss = gm.graph.call_method("type_as", args=(global_loss, input_node)) |
There was a problem hiding this comment.
Under torch.autocast with FP16/BF16 logits, PyTorch executes cross entropy in FP32 and returns an FP32 loss, but this unconditional type_as casts the aggregated scalar back to the low-precision logits. That changes the model's output contract and can round or overflow a loss before downstream loss composition or scaling; preserve the dtype recorded on the original loss node instead.
AGENTS.md reference: AGENTS.md:L38-L38
Useful? React with 👍 / 👎.
| """ | ||
| consumer = loss_node | ||
| target = _get_cross_entropy_arg(loss_node, "target", 1, None) | ||
| while _is_rank_preserving(target) or _is_reshape(target): |
There was a problem hiding this comment.
When the CE target is reshaped with sequence-dependent explicit dimensions, such as labels.view(batch, seq_len) or reshape(batch * seq_len), this loop accepts the reshape, but _localize_cross_entropy_target inserts an S/sp shard before it without rewriting its size arguments. Fake propagation or the first forward then tries to reshape the local target to the original full-sequence size and fails; restrict accepted target reshapes to inferred dimensions or rewrite the sequence-sized arguments.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b166cd8: a reshape on the target path with explicit symbolic sizes (e.g. labels.view(batch, seq_len)) is now rejected with an error pointing to view(-1), instead of failing later with an opaque shape error on the shorter local target. view(-1)/reshape(-1) remain supported. Covered by test_rejects_labels_it_cannot_align[explicit_view].
| weight = _get_cross_entropy_arg(loss_node, "weight", 2, None) | ||
| if weight is not None: | ||
| raise RuntimeError("AutoSP does not support class-weighted cross entropy on sequence-sharded labels") | ||
| reduction = _get_cross_entropy_arg(loss_node, "reduction", 6, "mean") |
There was a problem hiding this comment.
When a model supplies F.cross_entropy's still-supported size_average or reduce arguments, those arguments override reduction inside PyTorch. This code only reads reduction and later writes "none", so, for example, size_average=False is parsed as mean here while PyTorch keeps producing a summed scalar that AutoSP divides by the token count; reduce=False is likewise silently converted from a per-token result to a scalar mean instead of being rejected. Resolve the legacy arguments before choosing the aggregate mode.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in #8457 (e17b176): the deprecated size_average/reduce arguments are resolved with legacy_get_string (so e.g. size_average=False means sum and reduce=False is rejected as per-token), and cleared after the rewrite to reduction='none'. Covered by a size_average=False row in test_matches_full_sequence_reference and test_rejects_unsupported_cross_entropy.
- Raise when AutoSP is combined with ZeRO stage 2/3, which falls back to plain torch.compile while the SP mesh is already installed. - Reject class-probability targets, whose token count differs from the class-index path. - Honor legacy size_average/reduce arguments when resolving the reduction. - Cast the aggregated loss to the cross entropy output dtype instead of the logits dtype so autocast keeps the fp32 loss. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…odel label shift Keep the labels whole on every rank so that any shift the model applies to them, such as the Hugging Face F.pad(labels, (0, 1), value=-100)[..., 1:], runs on the full sequence. For each cross entropy on the labels, remove static sequence slices between the sharded logits and the loss (e.g. logits[..., :-1, :]), pad the full target to the global positions those slices covered, and hand the loss this rank's shard of the padded target. Only rank-preserving casts, reshapes after the slices, and static sequence slices are allowed between the logits and the cross entropy, and the labels may only flow into cross entropy targets. Any other pattern raises instead of silently pairing logits with the wrong targets. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
6f7c57b to
b166cd8
Compare
This branch has not been deployed
2 participants
Addresses #8457 (comment) (in-model causal label shift under AutoSP).
Problem
AutoSP shards the labels the same way as
input_ids. A causal LM that shifts its labels inside the model (Hugging FaceForCausalLMLosspads and shifts,labels = F.pad(labels, (0, 1), value=-100)[..., 1:]; Megatron-style code sliceslogits[..., :-1, :]/labels[..., 1:]) then shifts each shard on its own. Each rank pairs the last token of its shard with the wrong target, and the loss is silently wrong. #8457 documents pre-shifted labels as the workaround; this PR makes the in-model shift itself correct.Approach
F.cross_entropywhose target derives from the labels:logits[..., :-1, :]) are removed, because each rank must keep its whole shard of logits. Their offsets give the window[start, S + stop)of global positions the CE input covered.ignore_indexback to the global positions of that window, then hand the CE this rank's shard[r*S/sp, (r+1)*S/sp). This needs no communication because the labels are replicated.S + stop - start.contiguous,view/reshape/flatten(after any slices) and static sequence slices are allowed between the logits and the CE.input_ids, or returned from the module.narrow+full_like+catinstead ofF.pad/constant_pad_nd, which specialize the symbolic sequence length during fake-tensor propagation.Validation (8× Intel Arc Pro B70, torch 2.14.0.dev+xpu)
TestAutoSPCausalLosscompares the loss and the parameters after two optimizer steps against single-device training on the full sequence, over 4 ranks (sp=2, dp=2). New cases:slice(logits[..., :-1, :]/labels[..., 1:])pad(HF style)sum/mean, an explicit mesh, and an empty shardtest_rejects_labels_it_cannot_aligncovers a strided logits slice and labels masked byinput_ids.test_compile_autosp.py: 45 sequential and 20 non-sequential tests pass.TestAutoSPCompileandTestAutoSPEngineBackwardwere excluded because they fail on this machine for environmental reasons (dynamo tracing ofsynchronize(); the XPU C++20 op build, fixed in Use C++20 for XPU op builders on torch 2.12+ #8803), identically on the base branch.Tiny HF
LlamaForCausalLM(transformers 5.17), unshiftedlabels, sp=2, also checked withdynamic=True+attention_mask:Notes
shift_labels) remain supported and are still the simplest path.Contribution checklist
Related work. This PR is stacked on #8457 and addresses its review thread linked above. A search for "AutoSP label shift" found no other PR.
Sibling paths checked. AutoSP exists only as a DeepCompile pass, so there is no eager variant to fix. #8457 rejects ZeRO-2/3, and the label-shift handling here inherits that restriction.
Commands and results. Hardware and software: 4× Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu, transformers 5.17.0.
Results are listed under Validation above.
Numerics, performance, memory.
[B, S]integer ids per rank, negligible next to the logits.Compatibility matrix.
TestAutoSPCompilewas excluded, see above)AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.