执行摘要
- 一句话：修复嵌套张量重建对非最后 ragged 维的支持
- 推荐动作：该 PR 值得精读，展示了如何以最小改动修复一个因硬编码假设导致的通用性问题。读者可重点关注
nested_tensor_from_tensor_list 中 cat_dim 的推导以及 getattr 回退模式。同时建议跟进 review 中未解决的问题，尤其考虑添加 ragged_dim 参数到 nested_tensor_from_jagged 调用，并评估默认回退值是否应改为 1。
功能与动机
Issue #6152 报告，在在线策略蒸馏（OPD）等路径中，teacher_logprobs 张量的形状为 (batch, [seq_len], topk)，其 ragged 维度是 1（非最后）。PR #6127 引入的 nested_tensor_from_tensor_list 在 torch.cat 时硬编码 dim=-1，导致此类张量在 chunk/index_select 时崩溃。本 PR 修复此问题。
实现拆解
- 放宽 ragged_idx 断言：在
nested_tensor_from_tensor_list 中，将原先只允许 ragged_idx 等于最后一个维度的断言改为允许 1 ≤ ragged_idx ≤ sample_dim，并计算 cat_dim = ragged_idx - 1 用于 torch.cat。
- 传播 _ragged_idx 属性：在
concat_nested_tensors、chunk_tensordict、index_select_tensor_dict 中，从输入的嵌套张量读取 _ragged_idx 属性（若不存在则回退到 dim() - 1），并传递给 nested_tensor_from_tensor_list，确保重建后的张量继承正确的 ragged 维度。
- 添加回归测试：在
tests/test_protocol_v2_on_cpu.py 中新增两个测试，分别验证 chunk_tensordict 和 index_select_tensor_dict 在 ragged_idx=1 的 3D 嵌套张量上表现正确。
关键文件：
verl/utils/tensordict_utils.py（模块 数据工具；类别 source；类型 core-logic；符号 nested_tensor_from_tensor_list, concat_nested_tensors, chunk_tensordict, index_select_tensor_dict）: 核心修改文件，包含 nested_tensor_from_tensor_list、concat_nested_tensors、chunk_tensordict、index_select_tensor_dict 四个函数的改动，实现对任意 ragged 维度的支持。tests/test_protocol_v2_on_cpu.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_chunk_tensordict_preserves_3d_nested_tensor_layout_with_non_last_ragged_idx, test_index_select_tensor_dict_preserves_3d_nested_tensor_layout_with_non_last_ragged_idx）: 新增两个回归测试，验证非最后 ragged 维的嵌套张量在 chunk 和 index_select 操作中的正确性。
关键符号：nested_tensor_from_tensor_list, concat_nested_tensors, chunk_tensordict, index_select_tensor_dict
关键源码片段
verl/utils/tensordict_utils.py
核心修改文件，包含 nested_tensor_from_tensor_list、concat_nested_tensors、chunk_tensordict、index_select_tensor_dict 四个函数的改动，实现对任意 ragged 维度的支持。
def nested_tensor_from_tensor_list(
tensors: list[torch.Tensor], ragged_idx: int | None = None
) -> torch.Tensor:
"""从张量列表构建嵌套张量，支持指定任意 ragged 维度。
Args:
tensors: 张量列表，每个张量可以是 2D+（包含 batch 维度）。
ragged_idx: ragged 维度的索引（1-based），默认为最后一个维度。
Returns:
具有 jagged layout 的嵌套张量，其 `_ragged_idx` 属性被正确设置。
"""
assert len(tensors) > 0, "Must provide at least one tensor"
sample_dim = tensors[0].dim()
# ragged_idx 默认取最后一个维度（保持向后兼容）
if ragged_idx is None:
ragged_idx = sample_dim
# 放宽断言：允许 ragged 维度在 [1, sample_dim] 之间
assert 1 <= ragged_idx <= sample_dim, (
f"ragged_idx must be in [1, {sample_dim}]. Got {ragged_idx=} and {sample_dim=}"
)
if sample_dim == 1:
return torch.nested.as_nested_tensor(tensors, layout=torch.jagged)
# 根据 ragged_idx 计算 torch.cat 的维度（0-based）
cat_dim = ragged_idx - 1
# 沿 ragged 维度拼接 values
values = torch.cat(tensors, dim=cat_dim)
# 获取各样本在 ragged 维度的长度
lengths = torch.tensor(
[tensor.shape[cat_dim] for tensor in tensors],
dtype=torch.long,
device=values.device,
)
offsets = torch.zeros(len(tensors) + 1, dtype=torch.long, device=values.device)
torch.cumsum(lengths, dim=0, out=offsets[1:])
nested_tensor = torch.nested.nested_tensor_from_jagged(
values=values, offsets=offsets
)
# 设置 _ragged_idx 属性，供后续操作使用
nested_tensor._ragged_idx = ragged_idx
return nested_tensor
tests/test_protocol_v2_on_cpu.py
新增两个回归测试，验证非最后 ragged 维的嵌套张量在 chunk 和 index_select 操作中的正确性。
def test_chunk_tensordict_preserves_3d_nested_tensor_layout_with_non_last_ragged_idx():
"""回归测试：chunk_tensordict 必须能处理 ragged 维度不在最后一维的嵌套张量。"""
topk = 64
# 模拟 teacher_logprobs: (batch, [seq_len], topk)
elements = [
torch.randn(5, topk),
torch.randn(8, topk),
torch.randn(3, topk),
torch.randn(7, topk),
]
# 显式指定 ragged_idx=1，表示第一个维度（seq_len）是 ragged 维度
teacher_logprobs = tu.nested_tensor_from_tensor_list(elements, ragged_idx=1)
input_ids = torch.nested.as_nested_tensor(
[torch.arange(5), torch.arange(8), torch.arange(3), torch.arange(7)],
layout=torch.jagged,
)
td = tu.get_tensordict({"input_ids": input_ids, "teacher_logprobs": teacher_logprobs})
# 将 4 个 sample 等分为 2 组
chunks = tu.chunk_tensordict(td, chunks=2)
# 验证每个块的 _ragged_idx 得到保留
assert chunks[0]["teacher_logprobs"]._ragged_idx == 1
# 验证数据内容正确
assert torch.equal(chunks[0]["teacher_logprobs"].unbind(0)[0], elements[0])
assert torch.equal(chunks[0]["teacher_logprobs"].unbind(0)[1], elements[1])
assert chunks[1]["teacher_logprobs"]._ragged_idx == 1
assert torch.equal(chunks[1]["teacher_logprobs"].unbind(0)[0], elements[2])
assert torch.equal(chunks[1]["teacher_logprobs"].unbind(0)[1], elements[3])
评论区精华
Review 中 gemini-code-assist[bot] 提出了以下重点关注点：
风险与影响
- 风险：尽管 PR 解决了主要崩溃问题，仍存在风险：
- ragged_dim 参数未传递：
nested_tensor_from_jagged 默认 ragged_dim=1，若传入 ragged_idx=2 的 4D 张量，将错误地将第二维视为 ragged 维度，导致后续操作异常。
- 默认回退 dim()-1 不安全：对于标准 3D 嵌套张量（如
(batch, [seq_len], hidden)），若 _ragged_idx 属性缺失，回退到 dim()-1 将使用错误维度重建，可能引发静默数据损坏或运行时错误。
- 缺失 3D+ workaround：
index_select_tensor_dict 未实现 chunk_tensordict 中的 padded_tensor 回退路径，在遇到 3D+ 嵌套张量时仍可能崩溃。
- 测试覆盖不足：测试仅覆盖 ragged_idx=1 的情况，未覆盖 ragged_idx=2 等场景，也未覆盖
concat_nested_tensors 路径。
- 影响：
- 用户影响：修复了在线策略蒸馏（OPD）等路径中使用 teacher_logprobs 等 3D 嵌套张量时的崩溃问题，使相关训练流程恢复正常。
- 系统影响：修改了核心张量重建工具，提升了其对任意 ragged 维度的通用性，但也引入了上述风险。
- 团队影响：变更集中在
tensordict_utils.py，影响范围明确，但后续开发者需注意回退值的选择和参数传递的完整性。
- 风险标记：ragged_dim 参数未显式传递, 默认回退 dim()-1 不安全, 缺失 3D+ workaround, 测试覆盖不全
关联脉络
- PR #6127 [trainer] fix: preserve jagged tensor layout when rebuilding nested tensors with same sequence length: 引入了 nested_tensor_from_tensor_list 函数，其中硬编码 dim=-1 导致本 PR 修复的 bug。
- PR #6152 [trainer] bug: #6127 breaks distillation — nested_tensor_from_tensor_list assumes ragged dim is always the last: 报告了蒸馏路径中的崩溃问题，直接催生了本修复 PR。
参与讨论