Repository navigation
Conversation
![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: 0a392ba760
ℹ️ 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".
|
thanks for the PR, will take a closer look later, but can you help add more details of the issues, your fixes, and the test results. |
Sure, thank you. Will add more details asap. |
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>
|
Just to document the current status, @jinyouzhi clarified in #8667 that this PR is being split across multiple PRs. I also found a correctness issue with gradient averaging. Since I'm not sure how this PR overlaps with those, let's keep this PR open but hold off merging it for now. |
Thanks for verifying this patiently. There are still two issues to resolve. To avoid confusion, I’ll convert this PR to a draft. @tohtana |
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>
1578435 to
4806cd5
Compare
jinyouzhi
changed the title
@tohtana Thank you for your patience and for following up on this, I’ve refined #8457 and removed the parts that have already been merged (#8665, #8667, #8674, and #8675). The remaining AutoSP work is split across these PRs:
The gradient-averaging issue you raised is addressed in the latest #8457 revision. Since the engine averages gradients across all ranks, including SP ranks, the backward pass now applies an All four AutoSP PRs are still open. Please review only the last commit of #8457, which contains the causal-loss and gradient fix. |
… 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>
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Hi @jinyouzhi,
Thank you for the update! I left two comments about remaining correctness issues. For the shifted labels, I currently don't have a clear idea about the approach.
I think we can basically limit model/graph structures to ones and raise an explict error or warning when we are not sure the interpretation of the model is correct. Can you share your thought?
| continue | ||
| visited.add(node) | ||
| node_meta = get_node_shape_meta(node) | ||
| if isinstance(node_meta, torch.Tensor) and str(node_meta.shape) == label_shape: |
There was a problem hiding this comment.
This shape-based search still misses the conventional logits[..., :-1, :] / labels[..., 1:] pattern. The shifted target has length S-1, so the search walks back to the raw labels and inserts sharding before the shift, losing cross-boundary pairs such as C → D in A B C | D E F. This is an existing limitation that remains after the current fix.
I'm not sure what the most rubust and general approach would be. One approach is to use the semantics of the operations to determine where sharding is safe, rather than shape equality alone? Delaying label sharding might help, but the logits slice also needs to remove only the final position of the full sequence.
Do you see a good way to handle this pattern?
There was a problem hiding this comment.
Thanks, you're right that searching for the label by shape misses the [:-1] / [1:] shift. I've split this out of this PR:
- This PR ([AutoSP] Aggregate the causal LM loss and gradients over the SP group #8457) no longer tries to find the shifted label. It shards the raw labels, and the docs say labels must be shifted on the full sequence before
prepare_autosp_inputs()(e.g. passF.pad(labels, (0, 1), value=-100)[..., 1:]/shift_labels). - Draft [AutoSP] Align cross entropy targets after an in-model causal label shift #8811 (stacked on this PR) makes an in-model shift correct:
- Labels stay whole, so the shift runs on the full sequence.
- Static sequence slices between the sharded logits and the CE (
logits[..., :-1, :]) are removed. - Each CE gets this rank's shard of the target, padded back to global positions.
- Patterns it can't align raise an error instead of misaligning silently.
- On a tiny HF Llama with unshifted labels, it matches the single-device loss exactly ([AutoSP] Aggregate the causal LM loss and gradients over the SP group #8457 alone: 4.8458 vs 4.8407).
|
|
||
| loss = local_loss_sum.clone() | ||
| dist.all_reduce(loss, group=group) | ||
| # The engine averages gradients over every rank, SP ranks included, so the |
There was a problem hiding this comment.
The sp_size compensation fixes the earlier scaling issue on the implicit-mesh path. However, this assumption does not hold when an explicit DP×SP DeviceMesh is created via data_parallel_size or mesh_param: _get_sequence_data_parallel_group() falls back to the DP group, so the engine does not combine parameter-gradient contributions across SP ranks.
The reduction group and backward scaling seem to need to be handled together.
There was a problem hiding this comment.
Thanks, you were right: with data_parallel_size/mesh_param the sequence-data-parallel group fell back to the DP-only mesh dimension, so SP contributions were never summed (this also affected Ulysses with mesh_param). Fixed in fe38535 by making the engine own the SP reduction instead of compensating in the loss op:
groups._get_sequence_data_parallel_group()returns the whole (DP×SP) world when a mesh is set and there is nompu, so the engine sums over SP and averages over DP.deepspeed.initializecreates the same mesh when theautosppass is enabled with onlysequence_parallel_size, so implicit and explicit paths follow one convention.autosp::aggregate_lossno longer multiplies bysp_size.- Compiling AutoSP raises if an
mpu/mesh SP size differs fromsequence_parallel_size.
TestAutoSPCausalLoss now also runs with an explicit data_parallel_size mesh (ZeRO 0 mean, ZeRO 1 sum) against the single-device full-sequence reference, plus the mismatch case.
Split from deepspeedai#8457. ## Problem `prepare_autosp_inputs()` takes a `seq_dim` argument, but only passes it to `mark_dynamic`. The AutoSP graph passes ignore it: `pass_shard_seq_dim` and `shard_tensor_node` always read the sequence length from `shape[1]` and slice dimension 1. Sequence-first `[seq, batch]` inputs (`seq_dim=0`) are therefore mishandled: - If the batch dimension is static, compilation fails with `expected sequence dimension to be of type SymInt but found int`. - If the batch dimension is dynamic, AutoSP shards the batch instead of the sequence, and attention runs over the wrong tokens. ## Fix Dynamo copies only the `tag` attribute of an input tensor into graph metadata (`_extract_tensor_dict`), so `seq_dim` has to travel in the tag: - `prepare_autosp_inputs()` tags inputs as `(key, seq_dim)`. - `find_node_by_tag()` accepts both the tuple and a bare key. A bare key, as set manually in existing code and tests, keeps the previous default of dimension 1. - New `get_autosp_seq_dim()` reads the recorded dimension. `pass_shard_seq_dim` and `shard_tensor_node` use it to find the sequence length, and `shard_tensor_node` builds its slice for the tensor's full rank. - `label_id` is now bounds-checked against `seq_dim`, like the other prepared inputs, since it is also sharded along that dimension. - Docs: explain that `seq_dim` applies to every prepared tensor, and that `seq_dim=0` is for sequence-first inputs. ## Tests `TestAutoSPSeqDim` (2 ranks, `sp_size=2`) runs a small causal SDPA model with the sequence on `seq_dim` ∈ {0, 1}. For each of three batches, the mean of the per-rank AutoSP losses must match a single-device run on the full sequence (`rtol=1e-4`). Causal attention makes the loss depend on which dimension is sharded, so slicing the batch instead of the sequence would not match. - `seq_dim=0` fails on master and passes with this PR. - `seq_dim=1` passes on both, which covers the default path with the new tuple tags. The test compares forward losses only. AutoSP currently sums gradients over the SP group, while each rank's loss is a local mean, so gradients are `sp_size`× a full-sequence reference. That is part of the causal-loss aggregation follow-up from deepspeedai#8457, not this change. Results on 4× NVIDIA RTX 5090 D, torch 2.13.0, transformers 5.17.0: - `tests/unit/v1/compile/test_compile_autosp.py -k "not TestAutoSPCompile"`: 38 passed. - `TestAutoSPCompile` fails identically on unmodified master in this environment. The error is a Dynamo `InternalTorchDynamoError` in the Ulysses baseline's `torch.compile` forward, before AutoSP runs, so I excluded it. - `pre-commit` passes on the changed files. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: pengdurice <pengduhit@gmail.com>
Split out of deepspeedai#8457. ## Problem `get_node_shape_meta` in `deepspeed/compile/fx.py` picks between `meta["val"]` and `meta["example_value"]` with `or`. When `meta["val"]` is a tensor, that `or` evaluates the tensor's truthiness. As a result: - A tensor with more than one element raises `RuntimeError: Boolean value of Tensor with more than one value is ambiguous`. - A zero scalar is treated as missing, so the code falls back to a stale `example_value`. This only shows up after `pass_propagate_shapes`, which runs FakeTensorProp and fills in `meta["val"]`. Any AutoSP pass that reads shapes after that step can hit it. ## Fix Check `is not None` explicitly. ## Tests Added `TestNodeShapeMeta` to `tests/unit/v1/compile/test_compile_autosp.py`. It covers: - a multi-element tensor - a zero scalar - falling back to `example_value` Before the fix, 2 of these tests fail. After the fix, all 3 pass. Hardware: 8× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, transformers 5.17. The rest of the AutoSP unit tests pass. `TestAutoSPCompile` is not counted: it fails the same way on master in this environment, in the Dynamo Ulysses baseline. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ai#8761) Split out of deepspeedai#8457. **This PR is stacked on deepspeedai#8760.** It contains that PR's commit, so please review only the last commit. ## Problem AutoSP shards `input_ids`, `labels` and `position_ids` along the sequence dimension. A model can build its causal or attention mask from the sharded `position_ids`, which gives the mask the shape `[B, 1, S/sp, S]`. After the Ulysses all-to-all, SDPA sees all `S` queries, but the mask still covers only `S/sp` of them. The run then fails with a shape mismatch: `size of tensor a (s) must match ... (s//2)`. ## Fix - Add a custom op `autosp::all_gather_sequence(input, dim)` with a fake (meta) implementation. - In `pass_insert_attention_all_to_all`, check whether the mask's query dimension equals the local sequence length. If it does, all-gather the mask along dim -2 over the SP group. The mask can be passed as a kwarg or positionally. When several layers share one mask node, it is gathered only once. - If the mask's key dimension is also sharded, it cannot be rebuilt correctly from the local pieces, so raise a clear error. - Run `pass_propagate_shapes` once before the attention pass so that the mask shapes are current. This is why the PR depends on the `get_node_shape_meta` fix. - `all_gather_sequence` has a backward pass, an `autosp::reduce_scatter_sequence` op, so masks that require gradients (for example a learnable position bias) also work. - A mask with a full head dimension is redistributed like the query: if its query rows are sharded, it goes through all-to-all; otherwise each rank keeps only its own block of heads. A broadcast key dimension of 1 is never treated as sharded. - Document the behavior under AutoSP Limitations. ## Tests Added `TestAutoSPAttentionMask` (world_size=2): - With a learnable mask bias (shared, per-head with sharded queries, or per-head with full queries), the loss and all parameter gradients match a single-device reference. - With a query-sharded mask, passed as a kwarg and positionally, the forward loss matches a single-device full-sequence reference. - A mask whose key dimension is sharded is rejected with a clear error. Without the fix, all 3 tests fail. With the fix but without the `fx.py` change, all 3 also fail. Hardware: 8× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, transformers 5.17. 42 AutoSP tests pass. `TestAutoSPCompile` is excluded because of a pre-existing environment failure that also happens on master. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.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>
jinyouzhi
changed the title
- 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>
This branch has not been deployed
3 participants
Aggregates the AutoSP causal-LM loss and gradients over the sequence-parallel group, so training matches the full sequence.
This PR now covers only loss and gradient aggregation. Label alignment for a causal shift inside the model (tohtana's
sp_compile.pycomment) moves to a follow-up draft PR. Until that lands, labels must be shifted on the full sequence beforeprepare_autosp_inputs().Problem
AutoSP sharded the labels, and each rank computed a mean loss over only its own shard. This caused three problems:
sequence_parallel_sizeset): the engine sees no SP group and averages gradients over every rank.data_parallel_sizeormesh_param):groups._get_sequence_data_parallel_group()fell back to the DP-only mesh dimension, while the engine still divided bydp / sp, so SP contributions were never summed. This also affects Ulysses withmesh_param.Fix
autosp::aggregate_loss. EachF.cross_entropyon the labels is summed per shard in FP32. The op all-reduces the loss sum and the number of non-ignored tokens over the SP group, and every rank gets the full-sequence loss. Bothmeanandsumreductions work. In backward the gradient is scaled only by1 / total_validformean, with nosp_sizefactor.groups: summed SP gradients with a mesh. With a(data_parallel, sequence_parallel)mesh and nompu, the sequence-data-parallel group is the whole world. The engine then sums gradients over SP ranks and averages over DP ranks, which is the Ulysses convention.deepspeed.initialize: one convention for both paths. When theautosppass is enabled with onlysequence_parallel_size, it creates that mesh, so the implicit and explicit paths behave the same.mpuor mesh differs from AutoSP'ssequence_parallel_size.reduction="none"raise.Tests
TestAutoSPCausalLossruns with world_size=4 (sp=2, dp=2). The two shards hold different numbers of valid tokens. Over 2 SGD steps, it checks the loss and the updated parameters against single-device training on the full sequences, with gradient clipping off. Cases:meanandsumreductions, and keyword argumentsdata_parallel_sizeTestAutoSPHalfPrecisionLosscovers the FP16 shard-sum overflow. The attention-mask gradient tests were updated to the summed-gradient convention.Hardware: 8× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, transformers 5.17. All AutoSP tests pass except two groups that fail identically on master in this environment:
TestAutoSPCompile: dynamo fails tracingaccelerator.synchronize().TestAutoSPEngineBackward: the XPU op build needs C++20 (Use C++20 for XPU op builders on torch 2.12+ #8803).