参数

• context_parallel_config (ContextParallelConfig，可选) — 上下文并行的配置。
• tensor_parallel_config (TensorParallelConfig, 可选) — 张量并行的配置。

参数

• ring_degree (int, 可选，默认为 1) — 用于 Ring Attention 的设备数量。序列将在设备之间进行拆分。每个设备计算其本地 Q 与按环形顺序依次传递的 KV 块之间的注意力。显存占用更低（每次仅持有 1/N 的 KV），计算与通信可重叠，但需要 N 次迭代才能看到所有 token。最适合显存/带宽受限的长序列场景。在上下文并行区域内用于 ring attention 的设备数量。必须是上下文并行网格中设备总数的约数。
• ulysses_degree (int，可选，默认为 1) — 用于 Ulysses Attention 的设备数量。序列在设备之间进行切分。每个设备计算本地的 QKV，然后通过 all-gather 收集所有 KV 分块，以一次性计算完整的注意力。这种方式内存占用较高（需存储所有 KV），且需要高带宽的 all-to-all 通信，但延迟较低。最适合在互连带宽良好的情况下处理中等长度的序列。
• convert_to_fp32 (bool，可选，默认为 True) — 是否将输出和 LSE 转换为 float32，以确保 ring attention 的数值稳定性。
• rotate_method (str，可选，默认为 "allgather") — 在 ring attention 中用于跨设备轮换 key/value 状态的方法。目前仅支持 "allgather"。
• ulysses_anything (bool，可选，默认为 False) — 是否启用“Ulysses Anything”模式。该模式支持无法被 ulysses_degree 整除的任意序列长度和头数。启用时，ulysses_degree 必须大于 1，且 ring_degree 必须为 1。
• ring_anything (bool，可选，默认为 False) — 是否启用“Ring Anything”模式，该模式支持任意序列长度。启用时， ring_degree 必须大于 1，且 ulysses_degree 必须为 1。
• mesh (torch.distributed.device_mesh.DeviceMesh, 可选) — 用于上下文并行的自定义设备网格。如果提供，将使用该网格，而不是创建新的网格。当将上下文并行与其他共享同一设备网格的并行策略（例如 FSDP、张量并行）结合使用时，这非常有用。该网格必须同时具有 “ring” 和 “ulysses” 维度。对于未使用的维度，请使用大小 1（例如，仅使用 FSDP 的 ring attention 时，使用 mesh_shape=(2, 1, 4) 和 mesh_dim_names=("ring", "ulysses", "fsdp")）。

参数

• tp_degree (int，默认为 1) — 用于分片的设备数量。必须是正在并行化的模型的注意力头数（以及 FFN 隐藏维度）的约数。
• mesh (torch.distributed.device_mesh.DeviceMesh, 可选) — 要使用的自定义设备网格。如果提供，tp_degree 将从 mesh.size() 推断，该参数将被忽略。在将 TP 与其他共享同一网格的并行策略（例如 CP）结合使用时非常有用。

参数

• model (torch.nn.Module) — 要就地分片的模型。
• config (TensorParallelConfig) — 张量并行配置。其设备网格必须已通过 config.setup(...) 完成设置。
• tp_plan (dict) — 一个将模块名称通配符映射到 "colwise"/"rowwise" 样式（或其打包变体）的扁平映射，通常是模型的 _tp_plan。
• weights_already_sharded (bool，默认为 False) — 计划中的参数是否已经是 DTensor 分片，例如在通过流式 from_pretrained 加载之后。如果为 True，则仅注册前向钩子。此处显式传入该值而非自动检测，因为如果某个计划中的参数缺失于检查点中，它仍然会是一个 meta tensor，这会导致检测逻辑将一个实际上已部分分片的模型误判为“未分片”。