Stas Bekman

This commit introduces a **usability improvement** to error reporting within the **DeepSpeed Zero Stage 1 and 2 runtime**. Specifically, it **enhances an assertion message** in the `step` method of `deepspeed/runtime/zero/stage_1_and_2.py` that triggers when a parameter has already been reduced. Previously, this `AssertionError` would only display a generic numerical ID for the problematic parameter, making it difficult for users to identify the source of the issue. Now, the error message will explicitly state the **parameter's name**, such as `model.embed_tokens.weight`, providing much clearer and actionable debugging information for users encountering "double reduction" errors.

This commit provides a **Bug Fix** within the **DeepSpeed runtime engine** by correcting a **typo** in the `deepspeed/runtime/engine.py` module. The variable `savable_state_dict` is renamed to `saveable_state_dict` inside the `save_checkpoint` method. This ensures the **checkpoint saving mechanism** functions correctly, preventing potential errors or unexpected behavior when saving the model's state dictionary.

This commit **fixes a critical concurrency issue** within the **Modal CI workflows** by adjusting their GitHub Actions configuration. Previously, only a single Modal CI job could execute across all pull requests, leading to frequent cancellations of running jobs upon new PRs or updates. By modifying the concurrency group key in `.github/workflows/modal-accelerate.yml` and `.github/workflows/modal-torch-latest.yml` to utilize `github.ref`, this **CI/CD improvement** now enables multiple Modal CI jobs to run concurrently for different PRs. This change significantly **enhances CI efficiency** and prevents unnecessary job cancellations, ensuring more reliable and faster feedback for developers by optimizing resource utilization.

This commit **fixes** an issue within the **DeepSpeed Zero Stage 1 and 2** runtime where CPU memory was inadvertently being pinned even when the `cpu_offload_pin_memory` configuration was not explicitly enabled. A conditional check has been added to `deepspeed/runtime/zero/stage_1_and_2.py` to prevent this default behavior. This ensures that memory pinning for **CPU offloading** only occurs when explicitly configured by the user. The change improves memory management by avoiding unnecessary resource allocation and aligns the system's behavior with user-defined configurations for **DeepSpeed's memory optimization**.

This commit **improves the robustness** of the **DeepSpeed debugging utilities** by introducing helper functions `ds_id` and `ds_shape`. It **enhances resilience** within `deepspeed/utils/debug.py` by modifying existing debug functions, such as `debug_param2name_id` and `debug_param2name_id_shape`, to gracefully handle cases where `ds_*` attributes are not present on parameters. This **maintenance update** ensures that debugging operations continue without error, even when these specific attributes are missing, preventing potential crashes during introspection. The change primarily affects the **developer experience** by making debugging tools more reliable.

This commit **fixes** an inconsistency in the **memory usage utility** by replacing `logger.info` calls with `print` statements within the `see_memory_usage` function in `deepspeed/runtime/utils.py`. The previous implementation did not guarantee that memory usage messages would always be displayed, even when explicitly forced, which caused developer confusion and wasted debugging time. This **refactoring** ensures that critical memory information is **always printed** as intended, significantly improving **debugging reliability** and **developer experience** for the `deepspeed` runtime.

This commit **disables** the `nv-lightning-v100.yml` **CI workflow** within the project's GitHub Actions configuration. This **infrastructure maintenance** prevents the V100-based tests from running on pull requests, which was causing interference due to the unavailability of V100 GPUs. By commenting out the workflow's entire configuration, the project's **CI/CD pipeline** is streamlined, ensuring PRs are not blocked by failing V100 jobs. This is a temporary measure, with plans to port this testing functionality to Modal in the future.

This commit **improves code clarity** within the **DeepSpeed Zero Stage 3 runtime** by adding a detailed explanation. A new comment in `deepspeed/runtime/zero/stage3.py` clarifies the specific purpose of the gradient handling logic within the `reduce_leaf_module_grads` function. This **documentation enhancement** specifically addresses the rationale behind processing gradients for **Mixture of Experts (MoE) experts**. The change helps developers better understand the intricate gradient reduction mechanisms for large-scale models, serving as a **maintenance** update to improve code readability.

This commit **fixes** critical issues within the **`modal-accelerate` CI pipeline** to ensure reliable testing. It updates the `ci/accelerate.py` script to explicitly install `uv` in the CI image, a new requirement due to a base image change. Furthermore, it refactors the `accelerate` repository cloning and installation process, moving it into the `pytest` function to guarantee that the latest `accelerate` version is always used, preventing stale cached dependencies. Finally, the `.github/workflows/modal-accelerate.yml` workflow is updated with comments explaining how to properly test CI changes when using `pull_request_target`. This significantly improves the correctness and maintainability of the `modal-accelerate` integration tests.

This commit provides a **documentation update** for the **Ulysses Sequence Parallelism** implementation in DeepSpeed. It clarifies a critical performance characteristic of the `TiledMLP` and `SequenceTiledCompute` modules. The docstrings for these components in `deepspeed/runtime/sequence_parallel/ulysses_sp.py` now explicitly state that their memory-saving approach involves recomputing the forward pass during the backward pass. This ensures users are aware of the trade-off between memory efficiency and computational cost when utilizing these tiled MLP layers.

This commit **refactors the API** for **Ulysses Sequence Parallelism (SP)** to provide a more intuitive and flexible experience when dealing with **variable sequence lengths**. It **deprecates the `max_length` argument** in the `UlyssesSPAttentionHF.register_with_transformers` method, replacing it with `seq_length`, which is now optional if `seq_length_is_variable` is set to `True`. This **API improvement** primarily affects the `deepspeed/runtime/sequence_parallel/ulysses_sp.py` module, enhancing the integration with frameworks like Hugging Face Accelerate/Trainer. The change aims to make the API less confusing and more adaptable for dynamic input lengths. **Documentation and unit tests** have been updated to reflect these changes, ensuring clarity and correctness for users.

This commit **fixes a compatibility issue** within the **Ulysses (sequence parallel) MPU** implementation by adding necessary API aliases. It introduces `get_model_parallel_rank`, `get_model_parallel_world_size`, and `get_model_parallel_group` to `deepspeed.runtime.sequence_parallel.parallel_state_sp.py`, which now map directly to their sequence parallel counterparts. This **API enhancement** ensures that DeepSpeed's **checkpointing mechanism** can correctly interact with Ulysses, resolving an `AttributeError` that previously occurred when frameworks like **Hugging Face Trainer** attempted to save checkpoints. The change makes the Ulysses MPU fully compatible with existing DeepSpeed components that expect these model parallel APIs.

This commit **integrates Ulysses Sequence Parallelism (SP) with Hugging Face Accelerate**, enhancing the `UlyssesSPAttentionHF.register_with_transformers` method to directly accept a `PreTrainedModel` object, aligning with Accelerate's workflow. It also introduces a **defensive check** for sequence length within the `UlyssesSPDataLoaderAdapter` to improve robustness. Furthermore, unit tests for **Ulysses SP** and tiled compute are updated to test `zero_stage` 2 instead of 1, ensuring proper validation for this integration. This work provides a **new capability** for users leveraging Ulysses with Hugging Face Accelerate, improving compatibility and stability.

This commit **fixes a `KeyError`** that occurred in DeepSpeed's **Zero Stage 1 & 2 optimizer** during the gradient reduction process. Specifically, the `report_ipg_memory_usage` function would crash when `param.dtype` was passed as an argument, leading to a `KeyError` if the `ipg_buckets` dictionary did not contain an entry for that specific dtype, particularly with `fp32` communication settings. This **bug fix** resolves the issue by removing the `dtype` argument from the call to `report_ipg_memory_usage` within `reduce_independent_p_g_buckets_and_remove_grads` in `deepspeed/runtime/zero/stage_1_and_2.py`. This change aligns the behavior with the Zero Stage 3 implementation, preventing crashes and ensuring robust memory usage reporting when `seq_parallel_communication_data_type` is not `bf16`.

This commit **improves user experience** by enhancing an assertion message within the **DeepSpeed Zero Stage 1 and 2 runtime**. Specifically, the `reduce_ready_parameters` function in `deepspeed/runtime/zero/stage_1_and_2.py` will now report the actual parameter name (e.g., `lm_head.weight`) instead of a generic numeric ID when a parameter has already been reduced. This **debugging improvement** provides clearer context for users, making it significantly easier to identify and resolve issues related to parameter reduction.

This commit introduces **support for Mixture-of-Experts (MoE)** within the **TiledMLP** component, specifically addressing input shape handling. It **enhances** the `deepspeed/runtime/sequence_parallel/ulysses_sp.py` module to correctly process tensors when MoE routers drop the batch dimension. The `forward` and `backward` passes are adjusted to handle input shapes like `[seqlen, hidden_size]` by modifying tensor chunking and reshaping logic. This **new capability** ensures that **TiledMLP** can be seamlessly integrated and operate correctly within MoE-based model architectures, improving its versatility.

This commit performs a **refactoring** of **distributed logging utilities** within `deepspeed.utils.logging`. It introduces `get_dist_msg` to centralize common message formatting logic, then creates a new independent utility `print_dist` for direct distributed output, and updates the existing `log_dist` to leverage the new shared logic. This enhances code modularity and provides a dedicated function for printing distributed messages. As a direct impact, the `SynchronizedWallClockTimer` in `deepspeed.utils.timer` is updated to use `print_dist` for outputting timer statistics, ensuring consistent distributed reporting.

This commit **fixes a critical issue** where **wall clock breakdown statistics** were suppressed by high main logger levels, effectively disabling crucial performance insights. It introduces a `use_logger` parameter to the `log_dist` function in `deepspeed/utils/logging.py`, enabling direct printing of messages independent of the logger's level. Consequently, the `SynchronizedWallClockTimer` in `deepspeed/utils/timer.py` is updated to always print its statistics, ensuring this **performance monitoring feature** is consistently available. This **enhancement** improves the reliability and readability of performance debugging output by avoiding noisy logger prefixes.

This commit provides a **bug fix** for the **ZenFlow Adam optimizer** within `deepspeed.ops.adam`, specifically addressing an issue in `zenflow_torch_adam.py`. It ensures that the `_disable_dynamo_if_unsupported` fallback function is correctly defined, preventing a `NameError` during import when ZenFlow is unavailable. This resolves critical installation failures for **DeepSpeed on Torch 2.4 and 2.1**, which previously encountered import errors related to `torch.optim.optimizer`. Additionally, a redundant debug print statement has been removed as part of this **maintenance** effort, improving code cleanliness.

This commit **fixes** and **improves** the **Mixture of Experts (MoE) tutorial documentation**. It updates the URL for the CIFAR example, corrects existing text, and enhances the overall markup within the `docs/_tutorials/mixture-of-experts.md` file. This **documentation maintenance** ensures that users following the tutorial have access to accurate information and a better reading experience. The changes directly impact the usability of the **MoE feature's learning resources**, providing a more reliable guide for understanding and implementing MoE models.