ديون: مُحسِّن مايكروسوفت القابل للتطوير والفعال من حيث الاتصالات للتعلم الآلي الموزع

Dion, an optimizer developed by Microsoft, aims to address the communication bottlenecks present in distributed machine learning training, particularly when using orthonormal updates. Building upon the foundation laid by Muon, Dion introduces a more scalable approach for orthonormalizing weight matrices across multiple devices, enabling faster model convergence in large-scale training scenarios. This blog post will delve into the key features, implementation details, and usage guidelines of the Dion optimizer, based on the official documentation from the Microsoft GitHub repository.

Key Features of Dion

Dion offers a communication-efficient alternative to methods like Muon, particularly in scenarios involving sharded matrices. Its primary advantage stems from its use of amortized power iteration, a technique that allows for orthonormalization directly on sharded matrices, thereby reducing the need for communication-intensive matrix reconstruction. Key features include:

• Scalable Orthonormal Updates: Dion provides faster model convergence by computing orthonormal weight updates.
• Communication Efficiency: It reduces communication overhead by operating directly on sharded matrices, avoiding the need to reconstruct full matrices.
• Low-Rank Compression: Dion incorporates a rank fraction hyperparameter, enabling compute and communication reduction through low-rank approximation. An error feedback mechanism mitigates information loss.
• Parallelization Support: Dion supports single device, PyTorch DDP, and PyTorch FSDP2 parallelization. Notably, it is designed to work with FSDP2 + TP while Muon does not support this combination directly.
• Batch Processing and Overlapping: The optimizer processes parameters in batches and interleaves them to overlap compute with communication, improving performance.

Implementation and Usage

The Dion optimizer is available as a pip package, simplifying its integration into existing projects:

pip install git+https://github.com/microsoft/dion.git

Using Dion requires careful parameter grouping, as it’s designed specifically for two-dimensional matrix weights. Other parameters (biases, embeddings, normalizations) need separate optimizers like Lion or AdamW. Here are some key considerations for parameter grouping:

• Matrix Weights: Primarily for nn.Linear layers. Dion or Muon should be used, benefiting from an automatic shape-dependent learning rate scaling.
• Bias Vectors & Normalization Layers: Use Lion or AdamW.
• Embedding Layers: Treat them as a collection of 1D vectors and use Lion or AdamW (Dion will run without error, but perform poorly).
• Unembedding Layers (LM Head): Treat them as 1D vectors and apply a smaller scaled learning rate (critical for performance; use Lion or AdamW). Manual identification is essential.
• Convolution Layers: Experimental support in Muon using flatten=True. Dion does not support these.

Distributed Training Configuration

Proper configuration of the device mesh is crucial for Dion’s efficient operation in distributed environments. Dion utilizes DeviceMesh objects to understand the model’s parallelization scheme.

• Device Mesh for Dion: Supports up to two sharded mesh dimensions (outer_shard_mesh and inner_shard_mesh) and any number of data-parallel replicated mesh dimensions. Inner-shard-mesh benefits from intra-node tensor parallelism. Unused meshes can be omitted.
• Device Mesh for Muon: Uses a single 1D device mesh for either sharding parameters (efficient unsharding via all-to-all) or distributing work (all-gathering results). 2D sharding not supported - use Dion.
• DDP ProcessGroup: DDP is also supported, utilizing PyTorch’s ProcessGroup. If no process_group is given, Dion will redundantly execute the same work across all GPUs.

Compressed Data-Parallel Gradient Sync

Dion supports compressed data-parallel gradient synchronization, potentially skipping the full-gradient all-reduce by synchronizing only low-rank matrices, drastically reducing communication. replicate_mesh_grad_sync controls this feature:

• replicate_mesh_grad_sync=True: Dion all-reduces low-rank compressed states, resulting in decoupled momentum (needs explicit synchronization before checkpointing).
• replicate_mesh_grad_sync=False: Dion expects gradients have already been synchronized.

When using with HSDP (Hybrid Sharded Data Parallel), pass only the sharded sub-mesh to fully_shard() if Dion manages gradient sync.

Checkpointing

Before saving a checkpoint, optimizer.synchronize_for_checkpoint() must be called to synchronize optimizer states due to Dion’s decoupled momentum, ensuring consistency. DTensor types require saving checkpoints with torch.distributed.checkpoint.

Best Practices and Experimental Features

• Dion Rank Fraction: The most critical hyperparameter, controlling low-rank compression. Experiment to find the optimal value for your model size.
• Lion vs. AdamW: Lion generally performs better for scalar parameters when combined with Dion/Muon.
• 2D Sharding: When using both FSDP and TP, shard along different matrix dimensions.
• Learning Rate Scaling: Dion scales the learning rate automatically. Muon also supports the Moonshot AI scale factor.
• Experimental Features: Mixed precision and faster Dion through Cholesky QR (CQR) algorithm.

Conclusion

Dion offers a promising approach to scaling up distributed machine learning training by providing a communication-efficient optimizer that leverages orthonormal updates and low-rank compression. Understanding its specific requirements for parameter grouping, device mesh configuration, and checkpointing procedures is crucial for effectively utilizing Dion and achieving optimal performance in your distributed training workflows.

Source: Microsoft