LLM Systems Assignment 5: Distributed Training and Fine-Tuning

本次 homework 的任务是：实现 DP/PP/LoRA 的核心代码，整体代码量非常少，本文对实现要点和一些注意事项进行介绍。

1. Data Parallel

DP 部分的待实现代码非常简单，唯一需要提一点的是 ranndom 的相关使用方法：作业规定了不能使用普通的 random.shuffle 而是必须手动从指定 seed 构造 gen = random.Random(seed) 对象，然后调用 gen.shuffle(lst) 进行随机顺序打乱，这是因为：

random.shuffle 依赖于全局的 random state（seed），其是每个进程启动是随机初始化的，每个进程不同。
DP 在进行数据集划分时，需要保证全量数据集被不遗漏不重复的均匀划分到每个进程，partition_indices 函数使用单独的 seed 和 random stage，就是为了让每个进程计算得到的数据集 index shuffle 结果完全相同，然后每个进程取自己 rank 对应的 index 进行数据读取即可。

2. Pipeline Parallel

2.1 Clock-Cycle Schedule

以 3 microbatch、3 stage 为例，第 k 拍各 stage 在跑哪个 microbatch（行是时钟，元素写成 (microbatch, stage)）示意如下：

k	device 0	1	1
0	(0,0)
1	(1,0)	(0,1)
2	(2,0)	(1,1)	(0,2)
3		(2,1)	(1,2)
4			(2,2)

每个节拍满足不变量：第 k 拍恰好是所有 i + j == k（且 0<=i<microbatches、0<=j<stages）的 (i,j)，按 stage 递增排。总拍数 microbatches + stages - 1。

核心代码实现如下（省略参数校验）：

def clock_cycles(microbatches, stages):
    """Yield complete wavefronts of (microbatch, stage) pairs.

    Pair (i,j) runs at cycle i+j. Each cycle is ordered by increasing stage.
    Zero microbatches yields nothing. Reject negative microbatches or stages < 1.
    """

    for wave in range(microbatches + stages - 1):
        result = []
        for j in range(stages):
            micro_idx = wave - j
            if micro_idx >= 0 and micro_idx < microbatches:
                result.append((micro_idx, j))

        yield result

2.2 Pipeline Forward and Backward

2.2.1 StageWorker 接口

runtime.py 的 StageWorkers 是流水线执行引擎。每个 stage 一个后台线程，每个线程拥有两个队列（任务、结果）。

submit(stage, fn)：无参 callable 塞进任务队列，立刻返回，非阻塞。
receive(stage)：阻塞等结果（30s 超时）；worker 里抛的异常原样 raise。
close()：往队列塞 None 哨兵让 worker 退出。

几个直接影响写法的细节：

CUDA 当前 device 是线程局部的，worker 自己 with torch.cuda.device(...) 绑卡。所以提交的函数里算 tensor 不用每次显式 .to()。
梯度开关也是线程局部的。submit 时记录主线程此刻的 is_grad_enabled()，worker 执行前 set_grad_enabled 还原。
worker 用 try/except BaseException 把异常打包成 (False, exc)，receive 再 raise。所以 forward 里不用写 try/except。

2.2.2 Python lambda 特性

Python lambda 特性是：闭包晚绑定。lambda: model(micro_input.to(dev)) 里那几个名字，lambda 体在执行时（worker 线程，比定义晚）才读，那时循环早跑到最后一轮，所有 lambda 看到同一份最后的值。正确实现方式是在 submit 时把参数作为 xxx=xxx 显式初始化，这样参数立即求值构造为可调用对象提交到后台线程：

1	`lambda model=model, dev=dev, micro_input=micro_input: model(micro_input.to(dev))`

lambda 体只能是表达式，不能是语句。如果在 lambda 里填写： lambda ...: outputs[(m,s)] = model(x)，会报错：语法错误——赋值是 statement。正确实现方式是 lambda 只算并返回，结果在 receive 时保存。

跨设备 .to(dev) 一定写进被提交的 callable 里，让传输在 worker 线程、和各 stage 并行。

2.2.3 实现

一个 wave 循环，每拍先全 submit 再全 receive：

def forward(self, inputs):
    microbatches = torch.split(inputs, self.split_size)
    num_micro = len(microbatches)
    num_stage = len(self.stages)

    outputs = {}                       # (m, s) -> tensor

    for wave in clock_cycles(num_micro, num_stage):
        for micro_idx, stage_idx in wave:      # 先把整个 wave 提交出去
            model = self.stages[stage_idx]
            dev = self.devices[stage_idx]
            micro_input = (
                microbatches[micro_idx]
                if stage_idx == 0
                else outputs[(micro_idx, stage_idx - 1)]
            )
            self.workers.submit(
                stage_idx,
                lambda model=model, dev=dev, micro_input=micro_input: model(
                    micro_input.to(dev)
                ),
            )
        for micro_idx, stage_idx in wave:      # 再按同样顺序收结果
            outputs[micro_idx, stage_idx] = self.workers.receive(stage_idx)

    last_stage = num_stage - 1
    results = [outputs[micro_idx, last_stage] for micro_idx in range(num_micro)]
    return torch.cat(results)

数据依赖跨 wave。(m,s) 的输入是上一 stage 的 (m,s-1)，而 (m,s-1) 在上一拍（m+s-1）已经 receive 进 outputs。所以这次 submit 读 outputs[(m,s-1)] 时一定存在。outputs 字典在拍间传中间结果。
先提交完整 wave 再接收，以让各个 worker 并发执行。
一个 stage 一拍里最多出现一次，结果队列 FIFO，所以 receive 回来的就是当前 (m,s) 的结果，按与 submit 相同顺序 receive 即可获得结果。
不要 detach。micro_input.to(dev) 保留计算图，跨 stage 梯度链才连得上，bwd 才能把梯度传回输入和各 stage 参数。
拼接按 m 递增正好是原始顺序，最后在 self.devices[-1] 上 cat。

3. DeepSpeed ZeRO and LoRA

3.1 测试 DeepSpeed 初始化流程及与 TODO 代码交互关系

任务 3 的主要测试代码在：grading/distributed.py，其中涉及到 DeepSpeed 初始化以及调用我们需要实现的 TODO 代码。具体解释如下：

测试启动流程：
- main() 先用 socket.bind(("127.0.0.1", 0)) 检测一个空闲端口供 DeepSpeed 使用；进程组汇合用一个临时目录下的私有文件 store（init_method="file://..."）。
- mp.spawn(worker, nprocs=2) 起两个进程，把 port、store、case、accelerator 传进去。
- 每个 worker(rank, ...) 先写 MASTER_ADDR/PORT/RANK/LOCAL_RANK/WORLD_SIZE 环境变量，set_device(cuda:rank)，然后 dist.init_process_group("nccl", init_method="file://"+store, rank=rank, world_size=2, timeout=90s)。CPU 运行时后端是 gloo。
- torch.manual_seed(123) 固定随机种子，建 net = model() 和它的 ref 副本、loss_fn = MSELoss()。
与 TODO 代码交互关系：
- zero_runtime / zero_updates 分支：先 reference_adapters(net) 用 grader 自己的 LoRALinear 包住 net[0]、net[2]（net.requires_grad_(False)，所以只有 lora_A/lora_B 可训），ref = deepcopy(net) 作为单进程参考。optimizer = SGD(可训练参数)。
- 配置来源：zero_runtime 用 finetune.build_config(2, 2, 2, "fp32")，即调 TODO 的 build_config 函数。
- engine = deepspeed.initialize(model=net, optimizer=optimizer, config=config, dist_init_required=False) 把模型和 SGD 包成 DeepSpeedEngine（内部是 DeepSpeedZeroOptimizer）。
- 训练循环里，zero_updates 每个 microbatch 调 TODO 的 train_microbatch(engine, 切片, ...)。

3.2 模型 LoRA 权重替换：apply_lora

首先看 runtime.py 内已经提供的 LoRA 模块实现：

class LoRALinear(nn.Module):
    """Provided LoRA math; students implement targeting/freezing and persistence."""

    def __init__(self, base, rank, alpha):
        super().__init__()
        self.base = base
        self.scale = alpha / rank
        self.lora_A = nn.Parameter(base.weight.new_empty(rank, base.in_features))
        self.lora_B = nn.Parameter(base.weight.new_zeros(base.out_features, rank))
        nn.init.normal_(self.lora_A, std=0.02)
        self.base.requires_grad_(False)

    def forward(self, x):
        return self.base(x) + (x @ self.lora_A.T @ self.lora_B.T) * self.scale

其对应公式：

\[ h = W_0 x + b + \frac{\alpha}{r} B A x \] 其中 \(W_0\) 和 \(b\) 表示原始参数（freeze，不允许更新），\(A\) 和 \(B\) 表示可学习的 LoRA 参数。

apply_lora 函数的功能是：传入model 和需要替换的参数名列表 target_names 以及 LoRA 参数 rank 和 alpha，替换 model 内的需要替换的参数名对应的 Module 类型为上述的 LoRALinear ，从而实现模型的低参数量微调。

几个注意点：

替换 Linear Module 为 LoRALinear Module 方式：根据 parameter name 先获取 parent module，然后通过 setattr 设置 LoRALinear 替换原本模块。
冻结参数先用 model.requires_grad_(False) 让模型全部参数不可学习，再替换目标 Module 为 LoRALinear。

代码实现：

def apply_lora(model, target_names, rank=2, alpha=4):
    """Replace exactly the named nn.Linear modules with supplied LoRALinear.

    Validate all targets BEFORE mutating anything. Reject empty, duplicate,
    missing, or non-Linear targets and rank < 1 with ValueError. Freeze all base
    parameters; only lora_A and lora_B are trainable. Preserve device/dtype.
    Return the same model object; dotted nested names are supported.
    """
    # BEGIN_STUDENT
    named = dict(model.named_modules())

    if (
        len(target_names) == 0
        or rank < 1
        or len(set(target_names)) != len(target_names)
    ):
        raise ValueError("Invalid Arguments.")

    for name in target_names:
        if not name in named.keys() or not isinstance(named[name], nn.Linear):
            raise ValueError("Invalid Arguments.")

    model.requires_grad_(False)
    for name, module in model.named_modules():
        if name in target_names:
            parent_name, _, attr = name.rpartition(".")
            parent = model.get_submodule(parent_name) if parent_name else model
            setattr(parent, attr, LoRALinear(module, rank, alpha))
    return model
    # END_STUDENT

3.3 Save and Load Adapters

Save：提取模型中的 LoRA 相关参数，保存为 flatten dict(name:tensor)。
Load：加载 checkpoint flatten dict，拷贝从文件加载的 LoRA 参数到模型对应的 parameter。

代码实现如下：

def save_adapter(model, path):
    """Save CPU copies of ONLY named lora_A/lora_B tensors to path.

    Format is a flat name -> tensor dict loadable with weights_only=True.
    The harness calls this on a normal/gathered model, not partitioned ZeRO weights.
    Reject a model without adapters. Parent directory already exists.
    """
    # BEGIN_STUDENT
    adapter = {}
    for name, parameter in model.named_parameters():
        if name.endswith("lora_A") or name.endswith("lora_B"):
            adapter[name] = parameter.detach().cpu().clone()
    if len(adapter) == 0:
        raise ValueError("model without lora parameters")

    torch.save(adapter, path)
    # END_STUDENT


def load_adapter(model, path):
    """Restore a saved adapter into a model with matching adapters.

    Reject missing/extra keys or mismatched shapes before any parameter changes.
    Preserve the model's existing base weights and parameter registration.
    """
    # BEGIN_STUDENT
    saved = torch.load(path, weights_only=True)
    current = {
        name: p
        for name, p in model.named_parameters()
        if name.endswith(("lora_A", "lora_B"))
    }
    if set(saved.keys()) != set(current.keys()):
        raise ValueError("Invalid Arguments.")

    for name, parameter in current.items():
        if saved[name].shape != parameter.shape:
            raise ValueError("shape mismatch")

    with torch.no_grad():
        for name, value in saved.items():
            current[name].copy_(value)
    # END_STUDENT

4. Conclusion

本次 homework 主要收获点：

DP：
- 数据集切分注意事项
PP：
- wave 生成
- forward submit/receive 流程与中间 tensor 保存
LoRA
- LoRA 模块实现
- LoRA checkpoint 保存与加载

ML System

#CMU 11868 #LLM System #Distributed Training

LLM Systems Assignment 5: Distributed Training and Fine-Tuning

https://arcsin2.cloud/posts/2026/10/1893709119/

作者

arcsin2

发布于

2026年10月31日

许可协议

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 下一篇