Z1/2 init: flatten params on device #7828
Open
+155
−23
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ksugama
force-pushed
the
flatten-tensor-gpu
branch
from
a07a21b to
293fbab
Compare
ksugama
changed the title
Collaborator
|
@ksugama thanks for working on this. Looking forward to when it is ready for review. |
Fix deepspeedai#7812: This PR makes DeepSpeedEngine cleanup safe for partial initialization. This prevents destructor-time tracebacks by guarding access to unitialized attributes of DeepSpeed engine. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Evoformer tests fail with this error. We ignore this in the full test for now. ``` RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method ``` Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Two spelling errors in docs/_tutorials/accelerator-setup-guide.md: Line 50: comma-seperated-dash-range → comma-separated-dash-range Line 97: optimzied → optimized Both typos are in the Intel Architecture CPU section of the accelerator setup guide. Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
…er_rank (deepspeedai#7817) Fix deepspeedai#7811 Issue: ZeRO-3 crashes when `zero.GatheredParameters` is used with `modifier_rank=None` and a parameter is modified in-place. The failure surfaced as an internal invariant assert in free_param, which is not actionable for users. Solution: `GatheredParameters` now detects in-place mutation with `modifier_rank=None` and raises a clear, user-facing error early. The mutation check is synchronized across ranks to avoid divergence and hangs. This PR also raises a clearer error from free_param when parameters are still active in submodules.” --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
…rld_size in Ulysses (deepspeedai#7809) ### Description This PR addresses Issue deepspeedai#7672. When sequence_parallel_size is smaller than world_size (e.g., sp_size=2 on 4 GPUs) with PyTorch < 2.3, using torch.distributed.nn.functional.all_gather for loss aggregation triggers an IndexError: tuple index out of range during the backward pass. This is due to a known PyTorch issue where the backward hook accesses the global rank instead of the group rank. ### Solution 1. Regression Test & Workaround: Updated the regression test TestUlyssesLossBackward to implement a Weighted All-Reduce pattern. - Before: all_gather -> manual sum (Vulnerable to rank indexing mismatch on older PyTorch). - After: all_reduce(weighted_loss) / all_reduce(total_weight) (Robust and supports weighted averaging). 2. Runtime Warning: Added a version check (required_torch_version) in DeepSpeedEngine. It now logs a warning if Sequence Parallelism is enabled on PyTorch < 2.3, providing a link to the workaround test case. 3. Documentation: Updated ulysses-alst-sequence-parallelism.md with a note regarding legacy PyTorch versions and the recommended workaround. ### Verification Added and verified the regression test tests/unit/sequence_parallelism/test_ulysses.py which now validates the weighted averaging logic. **1. Reproduction (Before Fix)** Confirmed IndexError crash on Rank 2/3 with sp_size=2 on a 4-GPU setup. <img width="1370" height="860" alt="Screenshot 2026-01-23 at 23 53 42" src="https://github.com/user-attachments/assets/f4005c02-ff6c-46ea-a1a7-caac2093128b" /> **2. Verification (After Fix)** Verified the fix using the regression test logic on 4x RTX A6000. The backward pass now completes successfully on all ranks without error. <img width="1192" height="605" alt="Screenshot 2026-01-23 at 23 52 54" src="https://github.com/user-attachments/assets/c14cd093-67b7-42b0-ae15-65555c129082" /> --------- Signed-off-by: vensen <vensenmu@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Update PyTorch to v2.9 for modal tests --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Fix deepspeedai#7824 For [leaf modules](https://deepspeed.readthedocs.io/en/latest/training.html#configuring-zero-leaf-modules), ZeRO3 manages all parameters within the module uniformly. When a module returns multiple output tensors, PyTorch's autograd can trigger backward hooks from multiple threads concurrently. This causes race conditions when multiple threads simultaneously modify __inflight_param_registry and parameter states. This PR ensures that for leaf modules, only one thread performs the actual parameter fetching work while other concurrent threads wait and return early, preventing the race condition. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
hi deepspeed! [Make it possible to evaluate when using sequence parallel in HF Trainer #43517](huggingface/transformers#43517) I initially opened a PR in transformers to prevent errors that occur when running eval with deepspeed sequence parallel applied, but there was feedback that this should be handled in deepspeed rather than in transformers or accelerate, so I'm opening this PR here. If you have any questions, feel free to ask. Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
This PR introduces a flexible, configuration-driven API for AutoTP (Automatic Tensor Parallelism) that allows users to define custom layer partitioning patterns for training. @inkcherry @delock ## Motivation Previously, AutoTP relied on hardcoded layer detection logic that was difficult to customize for new model architectures. This PR enables: 1. **Custom models**: Users can define exact regex patterns to match their model's parameter names 2. **Fused layers**: Support for fused QKV, gate_up_proj, and other packed weight matrices with unequal sub-parameter sizes (e.g., GQA with different Q/K/V dimensions) 3. **Extensibility**: Easy to add new model presets or customize existing ones Here is an example of a config including custom partitioning patterns: ```json { "tensor_parallel": { "autotp_size": 4, "partition_config": { "use_default_specs": false, "layer_specs": [ { "patterns": [".*\\.o_proj\\.weight$", ".*\\.down_proj\\.weight$"], "partition_type": "row" }, { "patterns": [".*\\.[qkv]_proj\\.weight$"], "partition_type": "column" }, { "patterns": [".*\\.gate_up_proj\\.weight$"], "partition_type": "column", "shape": [2, -1], "partition_dim": 0 } ] } } } ``` Refer to the [document](https://github.com/tohtana/DeepSpeed/blob/tohtana/autotp_custom_patterns/docs/code-docs/source/training.rst) for more details (including preset models and how to define partitioning for fused models). We also opened a new [PR](deepspeedai/DeepSpeedExamples#998) to show the usage. ## Simplified initialization step AutoTP previously required calling ``set_autotp_mode(training=True)`` and ``deepspeed.tp_model_init`` before ``deepspeed.initialize``. Now we can include all the necessary configurations in the DeepSpeed config. We still support the traditional initialization path for backward compatibility. When you use both (i.e. calling ``set_autotp_mode(training=True)`` and ``deepspeed.tp_model_init`` and passing the config to ``deepspeed.initialize``), we will merge the settings at initialization. When we have conflicting settings, we will error out. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
Correctly handle `ds_grad_is_ready` in ZeRO2 --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
The current code has the following issues: - `use_default_specs: false` doesn't work - Injection by the traditional pattern runs even when custom patterns are set - `mpu` needs to be passed to `deepspeed.initialize` (HF integration doesn't pass mpu) This PR fixes AutoTP setup to respect `use_default_specs: false` and disable the traditional injection path when custom patterns are enabled. Also, when `mpu` is not passed, we create a TP group in the initialization process. With these changes, the [related tests](https://github.com/deepspeedai/DeepSpeed/tree/master/tests/unit/model_parallelism) pass and [all AutoTP examples](https://github.com/tohtana/DeepSpeedExamples/tree/tohtana/custom_auto_tp/training/tensor_parallel) in DeepSpeedExamples work now ([PR](deepspeedai/DeepSpeedExamples#998)). --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Kento Sugama <kentosugama@protonmail.ch>
ksugama
force-pushed
the
flatten-tensor-gpu
branch
from
bf577e8 to
3610631
Compare
ksugama
commented
Feb 9, 2026
| empty_cache() | ||
| see_memory_usage(f"After moving param group {i} to CPU", force=False) | ||
| accelerator = get_accelerator() | ||
| available_vram = accelerator.available_memory() if accelerator.is_available() else 0 |
Author
There was a problem hiding this comment.
@stas00 I believe available_memory() eventually calls into nvidia-smi. Is this the foot gun you were warning about?
If it is, maybe this should be fixed in a different PR since that problem touches more than is related to these changes
Author
|
@stas00 Hoping this alleviates the issue you were running into |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR addresses #7677 by flattening parameter tensors on the accelerators instead of the CPU during zero stage 1 and 2 initialization. This should alleviate CPU contention, with the caveat that the optimization is only used when there is enough VRAM to allocate a full copy of the parameter buffers.
On 8 x H100s and a Intel Xeon Platinum 8480+, profiling the initialization of DeepSpeed on 32 layers of
Qwen3-30Bwith Z2 gives the following:Old = ~382s
New = ~130s
If necessary, this optimization can be extended to allowed a tiered system that trades off VRAM space with performance, which might look like the following:
The distributed flatten would involve each device flattening a portion of the parameters and performing an all-gather to assemble the full flattened model. See #7677 for original discussion.