64 lines
3.9 KiB
Markdown
64 lines
3.9 KiB
Markdown
|
---
|
||
|
|
language:
|
||
|
|
- en
|
||
|
|
pipeline_tag: text-generation
|
||
|
|
tags:
|
||
|
|
- moe
|
||
|
|
- alignment
|
||
|
|
license: apache-2.0
|
||
|
|
library_name: transformers
|
||
|
|
---
|
||
|
|
|
||
|
|
# MoE-Pilot-Align-2.7B
|
||
|
|
|
||
|
|
## Overview
|
||
|
|
|
||
|
|
**MoE-Pilot-Align-2.7B** is a high-performance Sparse Mixture-of-Experts (MoE) model (14.3B total / 2.7B active), meticulously adapted and fine-tuned on a standard **8-GPU SXM node**.
|
||
|
|
|
||
|
|
This project is not a conventional fine-tuning exercise; it is an Infrastructure-Level Alignment showcase. It validates a production-ready pipeline for scaling MoE architectures by resolving critical bottlenecks in distributed initialization, heterogeneous memory management, and collective communication.
|
||
|
|
|
||
|
|
## The Engine: [Kernel-Align](https://github.com/Flink-ddd/Kernel-Align)
|
||
|
|
To achieve maximum throughput and memory efficiency during the alignment phase, this model is powered by **Kernel-Align**, an extreme post-training infrastructure optimized for NVIDIA and AMD GPUs.
|
||
|
|
|
||
|
|
* **Breaking the Memory Wall**: Optimized for **GRPO (Group Relative Policy Optimization)**, enabling a Group Size of **G=256** on a single A100 by keeping additional VRAM usage constant (**~0.5GB**).
|
||
|
|
* **Extreme Sampling Latency**: Integrated with **FlashInfer** and custom fused kernels, achieving up to **399x speedup** in the rollout phase compared to native PyTorch.
|
||
|
|
* **Universal Backend**: Native support for **AMD (ROCm/AITER)** and NVIDIA (CUDA), ensuring seamless cross-platform performance.
|
||
|
|
|
||
|
|
## Engineering & Infrastructure Innovations
|
||
|
|
|
||
|
|
### 1. Zero-Deadlock AOT Kernel Orchestration
|
||
|
|
In high-density 8-GPU environments, simultaneous JIT (Just-in-Time) compilation of MoE fused kernels often triggers distributed deadlocks due to I/O race conditions during the NCCL handshake.
|
||
|
|
- Innovation: Implemented an Ahead-of-Time compilation orchestrator. By decoupling CUDA/ROCm kernel building from the training lifecycle, we achieved zero-latency "Cold Starts" and eliminated synchronization timeouts.
|
||
|
|
|
||
|
|
### 2. Reflection-Based Parameter Metadata Injection
|
||
|
|
Standard DeepSpeed-ZeRO engines require precise expert-group tagging to avoid `AssertionError` during state partitioning. Hard-coded tagging is fragile and lacks portability.
|
||
|
|
- Solution: Developed a Reflection-based Metadata Injector. At the optimizer setup stage, the system dynamically scans tensors for `allreduce` attributes and naming patterns to inject `moe: True` tags.
|
||
|
|
- Impact: Ensured seamless compatibility between the MoE expert layers and the ZeRO-1/2 optimizer state sharding, maximizing VRAM efficiency on the 80GB A100 footprint.
|
||
|
|
|
||
|
|
### 3. Topology-Aware Parallelism (EP8)
|
||
|
|
- Parallelism Strategy: Optimized for a Single-Node 8-GPU Full-Mesh topology using **Expert Parallelism (EP=8)**.
|
||
|
|
- Communication Tuning: Refactored the `All-to-All` collective primitives to match the NVLink 3.0 bi-directional bandwidth. By balancing the load across 32 total experts (4 experts per GPU), we maintained a stable throughput of 185+ TFLOPs per device.
|
||
|
|
|
||
|
|
## Model Specifications
|
||
|
|
| Feature | Configuration |
|
||
|
|
| :--- | :--- |
|
||
|
|
| **Architecture** | Sparse MoE (Decoder-only) |
|
||
|
|
| **Total Parameters** | 14.3 Billion |
|
||
|
|
| **Active Parameters** | 2.7 Billion |
|
||
|
|
| **Experts Count** | 32 (Total) |
|
||
|
|
| **Routing Strategy** | Top-2 Gating with Load Balancing |
|
||
|
|
| **Precision** | Mixed Precision (BF16-O2) |
|
||
|
|
|
||
|
|
## Hardware & Software Environment
|
||
|
|
- **Compute**: 1 x Node | 8 x NVIDIA A100 80GB SXM
|
||
|
|
- **Interconnect**: NVLink 3.0 (600 GB/s Full-Mesh)
|
||
|
|
- **Framework**: Custom Fork of **Megatron-DeepSpeed**
|
||
|
|
- **CUDA/ROCm Compatibility**: Validated on CUDA 12.4; Architectural design supports seamless porting to AMD MI300X/ROCm environments via RCCL optimization.
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
## Community Quantizations
|
||
|
|
Special thanks to [@mradermacher](https://huggingface.co/mradermacher) for providing GGUF weights:
|
||
|
|
- [GGUF Weights (Q2_K to Q8_0)](https://huggingface.co/mradermacher/MoE-Pilot-Align-2.7B-GGUF)
|
||
|
|
-
|
||
|
|
*This repository serves as a technical benchmark for implementing production-scale MoE alignment protocols on next-generation heterogeneous compute clusters.*
|