DeepSpeed-Ulysses 实战：在 8×A100 上把 176B 模型压缩到 24GB 显存

# 安装命令示例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install deepspeed
pip install transformers

git clone https://github.com/microsoft/Megatron-DeepSpeed.git
cd Megatron-DeepSpeed
pip install -r requirements.txt

{
  "architectures": ["GPT2Model"],
  "hidden_size": 12288,
  "num_hidden_layers": 96,
  "num_attention_heads": 96,
  "max_position_embeddings": 32768,
  "vocab_size": 50257,
  "torch_dtype": "float16",
  "seq_parallel": true,
  "sequence_parallel_size": 8
}

{
  "train_batch_size": 512,
  "train_micro_batch_size_per_gpu": 1,
  "steps_per_print": 10,
  "zero_optimization": {
    "stage": 3,
    "contiguous_gradients": true,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6,
    "reduce_bucket_size": 1e8,
    "sub_group_size": 1e9,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },
  "fp16": {
    "enabled": true,
    "loss_scale_window": 100
  },
  "gradient_clipping": 1.0,
  "prescale_gradients": false,
  "wall_clock_breakdown": false,
  "sequence_parallel": {
    "enabled": true
  },
  "sparse_attention": {
    "mode": "fixed",
    "block": 16,
    "different_layout_per_head": true
  }
}

import os
import torch
import deepspeed
from megatron import get_args
from megatron.initialize import initialize_megatron
from megatron.model import GPTModel
from megatron.training import pretrain
from megatron.utils import get_ltor_masks_and_position_ids

# 必须通过命令行参数传递配置，这里在脚本中模拟核心逻辑
def model_provider():
    """构建模型"""
    args = get_args()
    model = GPTModel(
        num_tokentypes=0,
        parallel_output=True
    )
    return model

def get_batch(data_iterator):
    """生成一个虚拟批数据，实际应用需替换为真实数据加载器"""
    args = get_args()
    seq_length = args.seq_length
    batch_size = args.micro_batch_size * args.num_micro_batches
    
    # 构造随机数据 [batch, seq]
    tokens = torch.randint(0, args.padded_vocab_size, (batch_size, seq_length)).cuda()
    labels = torch.randint(0, args.padded_vocab_size, (batch_size, seq_length)).cuda()
    loss_mask = torch.ones(batch_size, seq_length, dtype=torch.float16).cuda()
    attention_mask = torch.tril(torch.ones((1, seq_length, seq_length))).expand(batch_size, -1, -1).cuda()
    
    return tokens, labels, loss_mask, attention_mask

def loss_func(loss_mask, output_tensor):
    """简化损失计算"""
    losses = output_tensor.float()
    loss_mask = loss_mask.view(-1).float()
    loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
    return loss

if __name__ == "__main__":
    # 在实际运行中，需要通过命令行传递所有参数。
    # 这里提供一个等效的命令行示例，用于启动训练：
    print("""
    实际启动命令示例（在 Megatron-DeepSpeed 目录下运行）：
    deepspeed --num_nodes=1 --num_gpus=8 train.py \
        --tensor-model-parallel-size 2 \
        --pipeline-model-parallel-size 2 \
        --sequence-parallel-size 2 \
        --num-layers 96 \
        --hidden-size 12288 \
        --num-attention-heads 96 \
        --seq-length 32768 \
        --max-position-embeddings 32768 \
        --batch-size 1 \
        --micro-batch-size 1 \
        --deepspeed \
        --deepspeed_config ds_config_ulysses.json \
        --data-path your_data_path \
        --tokenizer-type GPT2BPETokenizer \
        --vocab-file vocab.json \
        --merge-file merges.txt \
        --split 98,2,0 \
        --lr 1e-5 \
        --train-iters 1000 \
        --log-interval 1
    """)
    
    # 注意：上述命令中，并行度设置解释：
    # tensor-model-parallel-size (TP) = 2  # 张量并行维度
    # pipeline-model-parallel-size (PP) = 2 # 流水线并行维度
    # sequence-parallel-size (SP) = 2      # 序列并行维度（Ulysses）
    # 总 GPU 数 = TP * PP * SP = 2*2*2 = 8
    # 通过这种组合，我们实现了对参数、层和序列的多维度分解。

# 在训练脚本中，DeepSpeed会输出详细的显存使用情况
# 也可以通过以下命令动态观察
watch -n 1 nvidia-smi

[rank0] step=10, consumed samples=512, lr=1.00E-05, loss=7.34, ppl=1542.34
[rank0] memory_allocated (GB): 22.4 / 80.0
[rank0] max_memory_allocated (GB): 23.7 / 80.0