utils.py — TorchTitan 工具函数

文件概述

提供 TorchTitan 引擎的辅助工具（约 352 行），包括模型名推导、attention mask 构建、EP 权重收集等。

核心类和函数

NoOpDataLoader

class NoOpDataLoader(BaseDataLoader):
    """空操作数据加载器，因为 verl 自己管理数据加载"""
    def __iter__(self):
        return iter([])

derive_torchtitan_name_and_flavor

def derive_torchtitan_name_and_flavor(hf_config):
    """从 HuggingFace config 自动推导 TorchTitan 模型名

    映射关系:
    - qwen2/qwen3 → "qwen3"
    - llama → "llama3"
    - deepseek_v3 → "deepseek_v3"

    通过 hidden_size, num_hidden_layers, vocab_size 匹配具体 flavor
    """

enable_fsdp_gradient_division

def enable_fsdp_gradient_division(model, dp_size):
    """重新启用 FSDP 的梯度除法

    TorchTitan 默认禁用梯度除法（用于全局 token 归一化），
    但 verl 的损失函数需要梯度平均，所以这里重新启用。
    """

get_attention_masks

def get_attention_masks(input_batch, positions, attn_type):
    """构建 attention mask

    支持:
    - flex: FlexAttention（PyTorch 新特性）
    - varlen: 变长注意力（VarlenMetadata）
    """

iter_per_tensor_params_ep

def iter_per_tensor_params_ep(params, device, ep_group, ep_size):
    """EP 模式下的权重收集

    非专家参数: 直接 yield
    专家参数: 按 (layer, weight_type) 分组 all-gather，逐组 yield
    这样避免同时 materialize 所有专家权重导致 OOM
    """

与其他模块的关系

• 被 engine/torchtitan/transformer_impl.py 使用
• derive_torchtitan_name_and_flavor 是 HF → TorchTitan 模型映射的关键

小结

这些工具函数桥接了 verl 和 TorchTitan 框架之间的差异，特别是模型名映射和 EP 权重收集。