• 大模型
• 人工智能
• 多模态
• 媒体生成
• 基础模型

InternLumina-U2 是上海AI实验室推出的统一多模态模型：16B-A1B 的 MoE 稀疏主干 + 8 码本（codebook）全离散视觉表示（基于 AToken），一个模型覆盖文本问答、文生图、图像理解与编辑、视频与 3D 理解。训练这类模型时，最容易栽跟头的三个坑是FSDP2 集合通信挂起、多码本 Loss 的形状错配、序列 Packing 的注意力掩码。本文直接带读官方源码，把每类坑的根因和解法一次讲透，帮助你复现或微调时少烧卡。

模型速览：为什么它训练难

先看 nvidia/config.json 里几个关键参数，它们直接决定了训练会踩到哪些坑：

模型通过auto_map注册为LLaDA2MoeModel/LLaDA2MoeModelLM，核心实现在 nvidia/modeling_internluminau2.py；配置类在 nvidia/configuration_internluminau2.py。仓库中ascend/与nvidia/两个目录是同一架构分别在华为 Ascend NPU 和 NVIDIA GPU 上训练的 checkpoint，架构与权重布局一致（见 README.md）。

坑一：FSDP2 集合通信挂起与 NCCL 超时

现象：一个 rank 沉默，其余 rank 全部超时

训练突然卡住、无报错，几分钟后 NCCL watchdog 触发：多数 rank 报 collective timeout，而"肇事"的 rank 什么都不报——这基本就是 FSDP2 梯度集合通信不同步的典型特征。

根因：条件分支把参数甩出了 autograd 图

InternLumina-U2 的 forward 里有三类"条件性"输出头：图像答案头lm_img_head（8 个码本各一个 Linear）、文本答案头lm_head、图像码本头multi_codebook_head。图像分支只在label_mask.any()时执行，文本分支只在 pack 里有监督文本时执行。

一旦某个 rank 的分批恰好没有某一类样本（比如这个 pack 全是纯文本，没有图像 token），对应参数就不会进入该 rank 的 autograd 图，FSDP2 在该 rank 上就会少发起一次梯度集合通信，shard 组瞬间失步，整个作业挂死。源码注释写得很直白（modeling_internluminau2.py）：

A rank whose pack happens to carry none of one kind leaves those parameters out of the autograd graph. FSDP2 then issues one fewer gradient collective on that rank than on its peers, the shard group desyncs, and the job hangs until the NCCL watchdog fires.

解法：用零标量把所有条件头"钉"进计算图

项目采用_make_head_dummy()方案（modeling_internluminau2.py）：把所有条件模块的参数求和后乘以 0，加到 loss 里。零标量不影响 loss 数值，但能保证每个 rank 都对同一组参数参与梯度集合通信：

• _conditional_modules()：枚举全部条件模块——图像答案头、文本答案头、码本头、图像输入 embedding、图像投影、文本 embedding；
• 训练契约下forward会返回logits_img, logits_text, all_router_logits, None, self._make_head_dummy()（见 L2056-L2062），dummy 项负责"占位"。

💡 通用经验：凡是"按 batch 内容决定执行哪段网络"的 MoE/多模态模型，都要检查各 rank 的参数参与集是否恒定，否则 FSDP2 必挂。

附带细节：labels 的 dtype

FSDP2 混合精度可能把 kwargs 统一 cast 成 bf16，但 embedding 和交叉熵需要整型标签，源码因此显式检查并labels.long()转回（L1950-L1952）。如果你自定义数据管线，漏掉这一步会出现很难定位的 dtype 报错。

坑二：多码本 Loss 的标签形状与 -100 掩码

8 个码本 = 8 个独立的分类头

图像在 InternLumina-U2 里不是一条向量，而是每个位置一列8 个码本索引，因此答案头是一个 ModuleList——8 个Linear(hidden_size, 4096)分别预测每个码本（L1776-L1792）：

self.lm_img_head = nn.ModuleList([ nn.Linear(config.hidden_size, 4096, bias=False) for _ in range(self.num_codebooks) # 8 个码本各一个头 ])

中间还有一层MultiCodebookARHead（通道自回归，channel_depth=6），用前 7 个码本的 embedding 条件化第 8 个码本的预测，相关实现可参考 nvidia/tokenbridge.py。

三个最容易写错的点

1. labels 形状是[B, seq, C]，不是[B, seq]。Packing 场景下 labels 常见形状为[seq_length, num_codebooks]或[B, seq_length, num_codebooks]。源码先归一成 3D，再展平成[B*seq, C]，不能直接用按位置展开的img_ans_mask去索引已经展平过码本维度的张量（L1954-L1964）。
2. -100 掩码要跨码本取any。只要某位置在任意一个码本上有有效标签，该位置才参与图像分支 loss：label_mask = (labels != -100).any(dim=-1)。
3. 文本分支固定取第 0 个码本。文本位置每个位置只需要一个标签，源码对多码本 labels 统一text_labels[:, :, 0]取第 0 个码本的 token id（L2027-L2045）。

此外源码对valid image labels must be non-negative做显式校验，并在 embedding、hidden states、logits 多个节点插了torch.isfinite检查（抛FloatingPointError），比让 NaN 静默扩散到几十步后再炸要友好得多。

MoE 负载均衡 Loss 别忘加

256 个专家 + top-8 路由，负载均衡损失load_balancing_loss_func用经典的E * sum(f_i * P_i)形式（L134-L153），需要 forward 返回router_logits并加权进总 loss。仓库还内置了专家使用率追踪工具enable_expert_tracking/get_expert_stats（L87-L118），训练中途采样统计各层 top-k 命中数，能直接看到专家是否"偏科"。

坑三：序列 Packing 的块对角掩码与分段注意力

Packing 为什么必须带cu_seqlens

多模态 SFT 样本长度差异大，直接 padding 浪费算力。Packing 把多条样本拼成一条长序列，靠cu_seqlens（各段边界的累积长度）描述分段。backbone 收到cu_seqlens后构建严格块对角的注意力掩码：same_seg是一个硬 AND 项，且再叠加token_valid屏蔽 padding（L1372-L1408）。块扩散（block-diffusion）多轮样本还会叠加turn_ids/is_clean/block_ids：noisy 答案块可以看 clean 前文，clean token 永远不看 noisy 输出。

分段注意力：FLOPs 从 S² 降到 ΣLᵢ²

块对角掩码意味着 S×S 的稠密注意力可以逐段计算：use_segmented_attention开关打开时，走_segmented_sdpa，对每个 packed 段各调一次 SDPA kernel（L718-L750），算力从 S² 降到 ΣLᵢ²。几个工程细节值得抄作业：

• resolve_seg_bounds严格校验：cu_seqlens必须以 0 开头、严格递增、终点不超seq_len，错误直接抛异常而不是静默截断；且只在 backbone 每次 forward 转成 Python tuple一次，避免层循环里.tolist()带来的逐层 device sync（L683-L702）。
• 尾部 padding 的边界情况：全被 mask 的行用finfo.min而非-inf，会被平均成"对所有 key 的均匀平均"。为和稠密路径逐位对齐，tail pad 必须对着全量 K/V 做注意力（L734-L748）。
• SEG_ATTN_STRICT=1严格模式：config 声明分段注意力但实际走不了（缺cu_seqlens、处于增量解码、掩码不是 4D packing mask）时直接抛错，防止"checkpoint 里记录use_segmented_attention=true、实际却一直在跑 dense"这种静默漂移（L753-L782）。

别忽略 Packing collator 带来的 1D 张量

Packing 后attention_mask、position_ids常以 1D 形式传入（[S]），源码内部统一做unsqueeze(0)并对齐长度；position_ids始终是空间位置（B, seq_len），不带码本维度（L1351-L1368）。自己写 dataloader 时按这个契约输出即可，否则 rotary embedding 会拿到形状不对的输入。

专家并行（EP）怎么切：看 parallel_plan.py

16B-A1B 的 256 专家参数放不进单卡显存，模型用 FSDP2 的 DTensor 机制做专家并行。计划定义在 nvidia/parallel_plan.py：对每一层mlp.experts的gate_proj/up_proj/down_proj权重统一Shard(0)——即把 256 个专家沿第一维切到不同 rank。写自定义训练脚本时，这个 plan 必须被并行框架（veomni 的ParallelPlan）正确加载，否则专家切分方式不一致会直接 OOM 或算错。

上机前自查清单 ✅

1. 挂起先看 batch 构成：某 rank 的 pack 是否缺了图像或文本监督？条件头是否已用 head dummy 钉进计算图？
2. labels 三件事：3D 形状[B, S, C]先归一、-100跨码本取any、dtype 是整型。
3. cu_seqlens校验：以 0 开头、严格递增、终点 ≤ seq_len；严格模式下 tail 必须是纯 padding。
4. 负载均衡 loss：router_dtype=fp32下仍要记得把 aux loss 加权进总 loss。
5. EP plan 生效：专家投影权重Shard(0)，与训练框架的并行状态一致。

仓库结构速查

一句话总结：条件分支保参与集恒定、多码本标签先归一后掩码、Packing 靠cu_seqlens严格校验——把这三件事做对，InternLumina-U2 的多模态训练稳定性会提升一个量级。

• 大模型
• 人工智能
• 多模态
• 媒体生成
• 基础模型