Ring Attention 论文解析：在分布式环境中实现无限上下文窗口训练

import torch
import torch.nn.functional as F

def blockwise_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
                        block_size: int) -> torch.Tensor:
    """按块计算注意力（基于在线 Softmax）。

    不一次性计算完整的注意力分数矩阵，
    而是按块逐步累积出精确结果。

    Args:
        Q: Query 张量 [batch, seq_len, d_k]
        K: Key 张量 [batch, seq_len, d_k]
        V: Value 张量 [batch, seq_len, d_v]
        block_size: 块大小

    Returns:
        Output 张量 [batch, seq_len, d_v]
    """
    batch, seq_len, d_k = Q.shape
    d_v = V.shape[-1]
    scale = d_k ** -0.5
    num_blocks = seq_len // block_size

    output = torch.zeros(batch, seq_len, d_v, device=Q.device, dtype=Q.dtype)

    for i in range(num_blocks):
        q_block = Q[:, i * block_size:(i + 1) * block_size, :]  # [B, block_size, d_k]

        # 用于在线 Softmax 的累积变量
        max_score = torch.full((batch, block_size, 1), float('-inf'), device=Q.device)
        sum_exp = torch.zeros(batch, block_size, 1, device=Q.device)
        acc = torch.zeros(batch, block_size, d_v, device=Q.device)

        for j in range(num_blocks):
            k_block = K[:, j * block_size:(j + 1) * block_size, :]
            v_block = V[:, j * block_size:(j + 1) * block_size, :]

            # 计算部分注意力分数
            scores = torch.bmm(q_block, k_block.transpose(-2, -1)) * scale  # [B, bs, bs]

            # 在线 Softmax 更新
            new_max = torch.maximum(max_score, scores.max(dim=-1, keepdim=True).values)
            correction = torch.exp(max_score - new_max)
            new_exp = torch.exp(scores - new_max)

            # 修正此前的累积值并纳入新块
            sum_exp = sum_exp * correction + new_exp.sum(dim=-1, keepdim=True)
            acc = acc * correction + torch.bmm(new_exp, v_block)
            max_score = new_max

        # 最终归一化
        output[:, i * block_size:(i + 1) * block_size, :] = acc / sum_exp

    return output

class BlockwiseParallelTransformerLayer(torch.nn.Module):
    """BPT 层：按块融合处理注意力与 FFN。"""

    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)

        self.ffn = torch.nn.Sequential(
            torch.nn.Linear(d_model, d_ff),
            torch.nn.GELU(),
            torch.nn.Linear(d_ff, d_model),
        )
        self.norm1 = torch.nn.LayerNorm(d_model)
        self.norm2 = torch.nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor, block_size: int = 1024) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        num_blocks = seq_len // block_size

        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        output = torch.zeros_like(x)

        for i in range(num_blocks):
            start, end = i * block_size, (i + 1) * block_size
            q_block = Q[:, start:end, :]

            # 按块计算注意力（在线 Softmax）
            attn_out = self._blockwise_attn(q_block, K, V, block_size)
            attn_out = self.W_o(attn_out)

            # 残差连接 + 层归一化
            block_input = x[:, start:end, :]
            normed = self.norm1(block_input + attn_out)

            # 立即应用 FFN（节省内存的关键）
            ffn_out = self.ffn(normed)
            output[:, start:end, :] = self.norm2(normed + ffn_out)

        return output

import torch
import torch.distributed as dist

def ring_attention_forward(
    Q_local: torch.Tensor,   # 本设备的 Query 块 [batch, chunk_len, d_k]
    K_local: torch.Tensor,   # 本设备的 Key 块
    V_local: torch.Tensor,   # 本设备的 Value 块
    rank: int,               # 当前设备的 rank
    world_size: int,         # 设备总数
    scale: float,            # 缩放因子 1/sqrt(d_k)
) -> torch.Tensor:
    """Ring Attention 前向传播。

    在环形拓扑中循环传递 KV 块，按块累积计算注意力。
    以异步方式重叠通信与计算，隐藏通信开销。
    """
    batch, chunk_len, d_k = Q_local.shape
    d_v = V_local.shape[-1]

    # 在线 Softmax 累积变量
    max_score = torch.full((batch, chunk_len, 1), float('-inf'), device=Q_local.device)
    sum_exp = torch.zeros(batch, chunk_len, 1, device=Q_local.device)
    acc = torch.zeros(batch, chunk_len, d_v, device=Q_local.device)

    # 当前待处理的 KV 块（初始为本地块）
    kv_current = (K_local.clone(), V_local.clone())
    # 接收缓冲区
    kv_recv = (torch.empty_like(K_local), torch.empty_like(V_local))

    # 计算环上的邻居
    send_to = (rank + 1) % world_size
    recv_from = (rank - 1) % world_size

    for step in range(world_size):
        K_block, V_block = kv_current

        # 非最后一步时启动异步通信
        if step < world_size - 1:
            send_ops = [
                dist.isend(K_block, dst=send_to),
                dist.isend(V_block, dst=send_to),
            ]
            recv_ops = [
                dist.irecv(kv_recv[0], src=recv_from),
                dist.irecv(kv_recv[1], src=recv_from),
            ]

        # 对当前 KV 块计算部分注意力（与通信同时进行）
        scores = torch.bmm(Q_local, K_block.transpose(-2, -1)) * scale

        # 在线 Softmax 更新
        block_max = scores.max(dim=-1, keepdim=True).values
        new_max = torch.maximum(max_score, block_max)
        correction = torch.exp(max_score - new_max)
        new_exp = torch.exp(scores - new_max)

        sum_exp = sum_exp * correction + new_exp.sum(dim=-1, keepdim=True)
        acc = acc * correction + torch.bmm(new_exp, V_block)
        max_score = new_max

        # 等待通信完成后替换缓冲区
        if step < world_size - 1:
            for op in send_ops + recv_ops:
                op.wait()
            kv_current = (kv_recv[0].clone(), kv_recv[1].clone())

    # 最终归一化
    output = acc / sum_exp
    return output

def ring_attention_causal_step(
    Q_local: torch.Tensor,
    K_block: torch.Tensor,
    V_block: torch.Tensor,
    q_block_idx: int,      # Query 块在原始序列中的索引
    kv_block_idx: int,     # 当前 KV 块在原始序列中的索引
    block_size: int,
    scale: float,
) -> tuple:
    """应用了 causal masking 的 Ring Attention 单步计算。

    依据 KV 块的位置决定进行完整计算、部分掩码，还是完全跳过。
    """
    if kv_block_idx > q_block_idx:
        # KV 块的位置晚于 Query -> 完全跳过
        return None, None, None

    scores = torch.bmm(Q_local, K_block.transpose(-2, -1)) * scale

    if kv_block_idx == q_block_idx:
        # 仅在同一块内部应用部分 causal masking
        chunk_len = Q_local.shape[1]
        causal_mask = torch.triu(
            torch.ones(chunk_len, chunk_len, device=Q_local.device, dtype=torch.bool),
            diagonal=1
        )
        scores = scores.masked_fill(causal_mask.unsqueeze(0), float('-inf'))

    # kv_block_idx < q_block_idx：无需掩码（均为过去 token）

    block_max = scores.max(dim=-1, keepdim=True).values
    exp_scores = torch.exp(scores - block_max)
    block_sum = exp_scores.sum(dim=-1, keepdim=True)
    block_out = torch.bmm(exp_scores, V_block)

    return block_out, block_max, block_sum

def ring_attention_backward(
    dO_local: torch.Tensor,     # 输出梯度 [batch, chunk_len, d_v]
    Q_local: torch.Tensor,      # 保存下来的 Query
    K_local: torch.Tensor,      # 本地 Key
    V_local: torch.Tensor,      # 本地 Value
    O_local: torch.Tensor,      # Forward 输出
    lse_local: torch.Tensor,    # log-sum-exp（在线 Softmax 统计量）
    rank: int,
    world_size: int,
    scale: float,
) -> tuple:
    """Ring Attention 反向传播。

    以与 Forward 相同的 KV 循环模式计算 dQ、dK、dV。
    dK、dV 会累加回原本所属的设备。
    """
    batch, chunk_len, d_k = Q_local.shape
    d_v = V_local.shape[-1]

    dQ = torch.zeros_like(Q_local)
    dK_local = torch.zeros_like(K_local)
    dV_local = torch.zeros_like(V_local)

    # 预先计算 D 向量：rowsum(dO * O)
    D = (dO_local * O_local).sum(dim=-1, keepdim=True)  # [batch, chunk_len, 1]

    kv_current = (K_local.clone(), V_local.clone())
    send_to = (rank + 1) % world_size
    recv_from = (rank - 1) % world_size

    for step in range(world_size):
        K_block, V_block = kv_current

        # 重新计算注意力分数（由于使用了检查点，Forward 阶段未保存）
        scores = torch.bmm(Q_local, K_block.transpose(-2, -1)) * scale
        P = torch.exp(scores - lse_local)  # 归一化后的注意力权重

        # 计算梯度
        dV_block = torch.bmm(P.transpose(-2, -1), dO_local)
        dP = torch.bmm(dO_local, V_block.transpose(-2, -1))
        dS = P * (dP - D) * scale
        dQ += torch.bmm(dS, K_block)
        dK_block = torch.bmm(dS.transpose(-2, -1), Q_local)

        # 把 dK、dV 发送回原本所属的设备
        source_rank = (rank - step) % world_size
        if source_rank == rank:
            dK_local += dK_block
            dV_local += dV_block
        else:
            dist.reduce(dK_block, dst=source_rank, op=dist.ReduceOp.SUM)
            dist.reduce(dV_block, dst=source_rank, op=dist.ReduceOp.SUM)

        # 为下一步循环 KV
        if step < world_size - 1:
            kv_recv = (torch.empty_like(K_block), torch.empty_like(V_block))
            send_ops = [dist.isend(K_block, dst=send_to), dist.isend(V_block, dst=send_to)]
            recv_ops = [dist.irecv(kv_recv[0], src=recv_from), dist.irecv(kv_recv[1], src=recv_from)]
            for op in send_ops + recv_ops:
                op.wait()
            kv_current = kv_recv

    return dQ, dK_local, dV_local

import torch.distributed as dist

def hybrid_ulysses_ring_attention(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    intra_node_group: dist.ProcessGroup,  # 节点内组（Ulysses）
    inter_node_group: dist.ProcessGroup,  # 节点间组（Ring）
    n_heads: int,
) -> torch.Tensor:
    """USP：Ulysses + Ring Attention 混合序列并行化。

    节点内用基于 All-to-All 的 Ulysses 划分头维度，
    节点间用 Ring Attention 划分序列维度。
    """
    intra_size = dist.get_world_size(intra_node_group)
    inter_size = dist.get_world_size(inter_node_group)
    intra_rank = dist.get_rank(intra_node_group)
    inter_rank = dist.get_rank(inter_node_group)

    # Step 1：Ulysses All-to-All（节点内）
    # 序列维度划分 -> 重排为头维度划分
    Q_heads = all_to_all_reshape(Q, intra_node_group, split_dim='seq', gather_dim='head')
    K_heads = all_to_all_reshape(K, intra_node_group, split_dim='seq', gather_dim='head')
    V_heads = all_to_all_reshape(V, intra_node_group, split_dim='seq', gather_dim='head')

    # Step 2：Ring Attention（节点间）
    # 各节点上的设备针对各自的头执行 Ring Attention
    output_heads = ring_attention_forward(
        Q_heads, K_heads, V_heads,
        rank=inter_rank,
        world_size=inter_size,
        scale=(Q.shape[-1]) ** -0.5,
    )

    # Step 3：反向 All-to-All（节点内）
    # 头维度划分 -> 还原为序列维度划分
    output = all_to_all_reshape(output_heads, intra_node_group, split_dim='head', gather_dim='seq')

    return output

# 基于 JAX 的 Ring Attention 核心实现（参考官方代码）
import jax
import jax.numpy as jnp
from jax import lax

def ring_attention_jax(
    q: jnp.ndarray,    # [batch, chunk_len, n_heads, d_k]
    k: jnp.ndarray,    # [batch, chunk_len, n_heads, d_k]
    v: jnp.ndarray,    # [batch, chunk_len, n_heads, d_v]
    axis_name: str,     # pmap 轴名称
    scale: float,
    causal: bool = True,
    block_size: int = 1024,
) -> jnp.ndarray:
    """在 JAX pmap 环境下实现 Ring Attention。

    使用 lax.ppermute 让 KV 块在环上循环。
    """
    axis_size = lax.psum(1, axis_name)
    axis_index = lax.axis_index(axis_name)

    def scan_fn(carry, step):
        acc, max_score, sum_exp, k_block, v_block = carry

        # 当前 KV 块在原始序列中的索引
        kv_idx = (axis_index - step) % axis_size

        # Causal masking 检查
        if causal:
            should_compute = kv_idx <= axis_index
        else:
            should_compute = True

        # 计算部分注意力
        scores = jnp.einsum('bqhd,bkhd->bqhk', q, k_block) * scale

        if causal and kv_idx == axis_index:
            # 同一块：应用对角线 causal mask
            chunk_len = q.shape[1]
            mask = jnp.triu(jnp.ones((chunk_len, chunk_len)), k=1).astype(bool)
            scores = jnp.where(mask[None, :, None, :], -1e9, scores)

        # 在线 Softmax 更新
        new_max = jnp.maximum(max_score, scores.max(axis=-1, keepdims=True))
        correction = jnp.exp(max_score - new_max)
        new_exp = jnp.exp(scores - new_max)

        sum_exp = jnp.where(should_compute, sum_exp * correction + new_exp.sum(axis=-1, keepdims=True), sum_exp)
        acc = jnp.where(should_compute, acc * correction + jnp.einsum('bqhk,bkhd->bqhd', new_exp, v_block), acc)
        max_score = jnp.where(should_compute, new_max, max_score)

        # 把 KV 块循环传给环上的下一台设备
        # ppermute：按 (src, dst) 对交换数据
        perm = [(i, (i + 1) % axis_size) for i in range(axis_size)]
        k_block = lax.ppermute(k_block, axis_name, perm=perm)
        v_block = lax.ppermute(v_block, axis_name, perm=perm)

        return (acc, max_score, sum_exp, k_block, v_block), None

    # 初始化
    batch, chunk_len, n_heads, d_v = v.shape
    init_acc = jnp.zeros((batch, chunk_len, n_heads, d_v))
    init_max = jnp.full((batch, chunk_len, n_heads, 1), -1e9)
    init_sum = jnp.zeros((batch, chunk_len, n_heads, 1))

    init_carry = (init_acc, init_max, init_sum, k, v)
    (acc, max_score, sum_exp, _, _), _ = lax.scan(scan_fn, init_carry, jnp.arange(axis_size))

    return acc / sum_exp

import torch
from dataclasses import dataclass
from typing import List

@dataclass
class ContextSchedule:
    """渐进式上下文扩展的调度设置。"""
    context_lengths: List[int]     # 各阶段的上下文长度
    warmup_steps: List[int]        # 各阶段的训练步数
    rope_theta_values: List[float] # 各阶段的 RoPE theta 值

    def get_config(self, global_step: int) -> dict:
        cumulative = 0
        for i, steps in enumerate(self.warmup_steps):
            cumulative += steps
            if global_step < cumulative:
                return {
                    'context_length': self.context_lengths[i],
                    'rope_theta': self.rope_theta_values[i],
                    'stage': i,
                }
        return {
            'context_length': self.context_lengths[-1],
            'rope_theta': self.rope_theta_values[-1],
            'stage': len(self.context_lengths) - 1,
        }

# 4K -> 16K -> 64K -> 256K -> 1M 渐进式扩展示例
schedule = ContextSchedule(
    context_lengths=[4096, 16384, 65536, 262144, 1048576],
    warmup_steps=[1000, 800, 600, 400, 200],
    rope_theta_values=[10000, 50000, 500000, 5000000, 50000000],
)

# 在训练循环中使用
for step in range(3000):
    config = schedule.get_config(step)
    ctx_len = config['context_length']
    n_devices = torch.cuda.device_count()
    chunk_per_device = ctx_len // n_devices

    print(f"Step {step}: context={ctx_len}, "
          f"chunk/device={chunk_per_device}, "
          f"RoPE theta={config['rope_theta']:.0f}, "
          f"stage={config['stage']}")

from typing import List, Optional
import torch

class LongContextDataCollator:
    """用于 Ring Attention 训练的长文本数据 Collator。

    把多份文档连接起来构成目标序列长度，
    并做填充使其能均匀分配给 N 台设备。
    """

    def __init__(
        self,
        tokenizer,
        target_seq_len: int,
        world_size: int,
        doc_separator_id: int = 2,  # </s> 或 <|endoftext|>
    ):
        self.tokenizer = tokenizer
        self.target_seq_len = target_seq_len
        self.world_size = world_size
        self.doc_separator_id = doc_separator_id
        self.chunk_size = target_seq_len // world_size

    def __call__(self, documents: List[str]) -> dict:
        # 对文档分词并用分隔符连接
        all_tokens = []
        doc_boundaries = []
        for doc in documents:
            tokens = self.tokenizer.encode(doc, add_special_tokens=False)
            doc_boundaries.append(len(all_tokens))
            all_tokens.extend(tokens)
            all_tokens.append(self.doc_separator_id)

        # 按目标长度截断或填充
        if len(all_tokens) > self.target_seq_len:
            all_tokens = all_tokens[:self.target_seq_len]
        elif len(all_tokens) < self.target_seq_len:
            pad_len = self.target_seq_len - len(all_tokens)
            all_tokens.extend([self.tokenizer.pad_token_id] * pad_len)

        # 确认能否被 world_size 整除
        assert len(all_tokens) % self.world_size == 0, (
            f"시퀀스 길이 {len(all_tokens)}가 "
            f"world_size {self.world_size}로 나눠지지 않습니다."
        )

        input_ids = torch.tensor(all_tokens, dtype=torch.long)

        # 生成分配给各设备的 chunk 索引
        chunks = input_ids.view(self.world_size, self.chunk_size)

        return {
            'input_ids': input_ids,
            'chunks': chunks,
            'doc_boundaries': doc_boundaries,
        }