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:
    """Blockwise 어텐션 계산 (Online 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]

        # 온라인 소프트맥스를 위한 누적 변수
        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]

            # 온라인 소프트맥스 업데이트
            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, :]

            # 블록별 어텐션 (온라인 소프트맥스)
            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,               # 현재 디바이스 순위
    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]

    # 온라인 소프트맥스 누적 변수
    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

        # 온라인 소프트맥스 업데이트
        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 블록의 위치에 따라 Full 계산, 부분 마스킹, 완전 스킵을 결정한다.
    """
    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: 마스킹 불필요 (모두 과거 토큰)

    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,     # 출력 gradient [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 (온라인 소프트맥스 통계)
    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)  # 정규화된 어텐션 가중치

        # Gradient 계산
        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)

        # 온라인 소프트맥스 업데이트
        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)

        # 각 디바이스에 할당될 청크 인덱스 생성
        chunks = input_ids.view(self.world_size, self.chunk_size)

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