sp_megatron.py

import os

os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1"
os.environ["NCCL_DEBUG"] = "WARN"

import copy

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
import torch.distributed.tensor.parallel as tp
from torch.utils.benchmark import Timer

import peft
from transformers.models.llama.modeling_llama import LlamaAttention, LlamaMLP

from utils import setup, load_model_on_rank0, load_model_serialized
from config import DEFAULT_PROMPT


class ComputeWithAllGather(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tp_shard: nn.Module, input: torch.Tensor, args, kwargs):
        world_size = dist.get_world_size()
        gathered_input = [torch.empty_like(input) for _ in range(world_size)]
        dist.all_gather(gathered_input, input)
        gathered_input = torch.cat(gathered_input, dim=1)
        if "position_ids" in kwargs:
            pos = kwargs["position_ids"]
            gathered_pos = [torch.empty_like(pos) for _ in range(world_size)]
            dist.all_gather(gathered_pos, pos)
            kwargs["position_ids"] = torch.cat(gathered_pos, dim=1)
        if "position_embeddings" in kwargs:
            cos, sin = kwargs["position_embeddings"]
            gathered_cos = [torch.empty_like(cos) for _ in range(world_size)]
            gathered_sin = [torch.empty_like(sin) for _ in range(world_size)]
            dist.all_gather(gathered_cos, cos)
            dist.all_gather(gathered_sin, sin)
            kwargs["position_embeddings"] = (torch.cat(gathered_cos, dim=1), torch.cat(gathered_sin, dim=1))
        if "attention_mask" in kwargs and kwargs["attention_mask"] is not None:
            seq_len = gathered_input.shape[1]
            causal = torch.triu(
                torch.full((seq_len, seq_len), torch.finfo(gathered_input.dtype).min, device=gathered_input.device),
                diagonal=1,
            )
            kwargs["attention_mask"] = causal.unsqueeze(0).unsqueeze(0)
        # Save only the local shard, not the gathered tensor
        ctx.save_for_backward(input)
        ctx._tp_shard = tp_shard
        ctx._args = args
        ctx._kwargs = kwargs
        output = tp_shard(gathered_input, *args, **kwargs)
        out0 = output[0] if isinstance(output, tuple) else output
        chunks = list(out0.chunk(world_size, dim=1))
        scattered = torch.empty_like(chunks[0])
        dist.reduce_scatter(scattered, chunks)
        if isinstance(output, tuple):
            return (scattered, *output[1:])
        return scattered

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor, *rest_grads):
        (input_shard,) = ctx.saved_tensors
        world_size = dist.get_world_size()
        # re-gather input from shards
        gathered_input = [torch.empty_like(input_shard) for _ in range(world_size)]
        dist.all_gather(gathered_input, input_shard)
        gathered_input = torch.cat(gathered_input, dim=1).detach().requires_grad_(True)
        # gather grad_output
        gathered_grad_output = [torch.empty_like(grad_output) for _ in range(world_size)]
        dist.all_gather(gathered_grad_output, grad_output)
        gathered_grad_output = torch.cat(gathered_grad_output, dim=1)
        with torch.enable_grad():
            output = ctx._tp_shard(gathered_input, *ctx._args, **ctx._kwargs)
            out0 = output[0] if isinstance(output, tuple) else output
            out0.backward(gathered_grad_output)
        gathered_input_grad = list(gathered_input.grad.chunk(world_size, dim=1))
        input_grad = torch.empty_like(gathered_input_grad[0])
        dist.reduce_scatter(input_grad, gathered_input_grad)
        return None, input_grad, None, None


class AllGatherModule(nn.Module):
    def __init__(self, module):
        super().__init__()
        self.module = module

    def forward(self, *args, **kwargs):
        if args:
            input = args[0]
            rest_args = args[1:]
        else:
            input = kwargs.pop("hidden_states")
            rest_args = ()
        return ComputeWithAllGather.apply(self.module, input, rest_args, kwargs)


class SPModelWrapper(nn.Module):
    def __init__(self, model, gather_output=False):
        super().__init__()
        self.model = model
        self.gather_output = gather_output

    def _scatter(self, tensor):
        return tensor.chunk(dist.get_world_size(), dim=1)[dist.get_rank()]

    def _position_ids(self, seq_len, batch_size, device):
        ws = dist.get_world_size()
        cs = seq_len // ws
        return (
            torch.arange(dist.get_rank() * cs, (dist.get_rank() + 1) * cs, device=device)
            .unsqueeze(0)
            .expand(batch_size, -1)
        )

    def forward(self, input_ids=None, inputs_embeds=None, labels=None, attention_mask=None, **kwargs):
        if input_ids is not None:
            input_ids = self._scatter(input_ids)
        if inputs_embeds is not None:
            inputs_embeds = self._scatter(inputs_embeds)
        if labels is not None:
            labels = self._scatter(labels)
        if attention_mask is not None:
            attention_mask = self._scatter(attention_mask)

        # we need global position_ids (note - they will be ignored in prompt tuning)
        seq_len = (input_ids if input_ids is not None else inputs_embeds).shape[1]
        batch_size = (input_ids if input_ids is not None else inputs_embeds).shape[0]
        device = (input_ids if input_ids is not None else inputs_embeds).device
        position_ids = self._position_ids(seq_len * dist.get_world_size(), batch_size, device)

        output = self.model(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            labels=labels,
            attention_mask=attention_mask,
            position_ids=position_ids,
            **kwargs,
        )

        if self.gather_output:
            # it is hack to make autograd work
            all_logits = [torch.zeros_like(output.logits) for _ in range(dist.get_world_size())]
            dist.all_gather(all_logits, output.logits)
            all_logits[dist.get_rank()] = output.logits
            output.logits = torch.cat(all_logits, dim=1)
        return output


def make_sequence_parallel(model, ctx):
    rank, world_size, config = ctx.rank, ctx.world_size, ctx.config
    l = list(model.named_children())
    for name, child in l:
        tp_config = copy.deepcopy(config)
        if isinstance(child, LlamaAttention):
            tp_config.num_attention_heads = config.num_attention_heads // world_size
            tp_config.num_key_value_heads = config.num_key_value_heads // world_size
            tp_attn = LlamaAttention(tp_config, layer_idx=child.layer_idx).to(child.q_proj.weight.dtype)
            with torch.no_grad():
                tp_attn.q_proj.weight.copy_(torch.chunk(child.q_proj.weight, world_size, dim=0)[rank])
                tp_attn.k_proj.weight.copy_(torch.chunk(child.k_proj.weight, world_size, dim=0)[rank])
                tp_attn.v_proj.weight.copy_(torch.chunk(child.v_proj.weight, world_size, dim=0)[rank])
                tp_attn.o_proj.weight.copy_(torch.chunk(child.o_proj.weight, world_size, dim=1)[rank])
            setattr(model, name, AllGatherModule(tp_attn))
        elif isinstance(child, LlamaMLP):
            tp_config.intermediate_size = config.intermediate_size // world_size
            tp_mlp = LlamaMLP(tp_config).to(child.gate_proj.weight.dtype)
            with torch.no_grad():
                tp_mlp.gate_proj.weight.copy_(torch.chunk(child.gate_proj.weight, world_size, dim=0)[rank])
                tp_mlp.up_proj.weight.copy_(torch.chunk(child.up_proj.weight, world_size, dim=0)[rank])
                tp_mlp.down_proj.weight.copy_(torch.chunk(child.down_proj.weight, world_size, dim=1)[rank])
            setattr(model, name, AllGatherModule(tp_mlp))
        else:
            make_sequence_parallel(child, ctx)
    return model


def test_correctness(sp_llama, ref_llama, prompt, ctx):
    input_ids = ctx.tokenizer(prompt, return_tensors="pt")["input_ids"]
    pad_len = (ctx.world_size - input_ids.shape[1] % ctx.world_size) % ctx.world_size
    if pad_len > 0:
        input_ids = F.pad(input_ids, (pad_len, 0), value=ctx.tokenizer.pad_token_id or 0)

    sp_embeds = sp_llama.model.model.embed_tokens(input_ids.to(ctx.device)).detach().requires_grad_(True)
    sp_logits = sp_llama(inputs_embeds=sp_embeds, use_cache=False).logits
    sp_logits.mean().backward()

    my_chunk = sp_embeds.grad.chunk(ctx.world_size, dim=1)[ctx.rank]
    chunks = [torch.empty_like(my_chunk) for _ in range(ctx.world_size)]
    dist.all_gather(chunks, my_chunk)
    sp_grad = torch.cat(chunks, dim=1)

    if ctx.rank == 0:
        ref_embeds = ref_llama.model.embed_tokens(input_ids).detach().requires_grad_(True)
        ref_logits = ref_llama(inputs_embeds=ref_embeds, use_cache=False).logits
        ref_logits.mean().backward()

        logit_diff = (sp_logits.cpu() - ref_logits).detach().abs().max()
        grad_diff = (sp_grad.cpu() - ref_embeds.grad).abs().max()
        print(f"logits diff: {logit_diff:.6f} {'OK' if logit_diff <= 2e-5 else 'MISMATCH'}")
        print(f"grads diff:  {grad_diff:.6f} {'OK' if grad_diff <= 2e-5 else 'MISMATCH'}")


def test_learning(sp_llama, ctx):
    for p in sp_llama.parameters():
        p.requires_grad_(False)

    torch.manual_seed(42)
    torch.cuda.manual_seed(42)
    input_ids = torch.randint(0, ctx.config.vocab_size, (1, ctx.length), device=ctx.device)

    with torch.no_grad():
        embeds = sp_llama.model.model.embed_tokens(input_ids)
    embeds = embeds.detach().clone().requires_grad_(True)

    labels = input_ids.clone()
    opt = torch.optim.Adam([embeds], lr=1e-3)

    for i in range(5):
        local_loss = sp_llama(inputs_embeds=embeds, labels=labels, use_cache=False).loss
        opt.zero_grad()
        local_loss.backward()
        opt.step()

        with torch.no_grad():
            loss = local_loss.clone()
            dist.all_reduce(loss)
            loss /= ctx.world_size

        if ctx.rank == 0:
            print(f"{i=}\t{loss.item()=}")


def benchmark(sp_llama, ctx):
    input_ids = labels = torch.randint(0, ctx.config.vocab_size, (1, ctx.length), device=ctx.device)

    def fwd():
        return sp_llama(input_ids=input_ids, use_cache=False).logits

    def fwd_bwd():
        sp_llama(input_ids=input_ids, use_cache=False).logits.mean().backward()

    fwd_t = Timer(stmt="fwd()", globals={"fwd": fwd}).timeit(ctx.bench_iters)
    fwd_bwd_t = Timer(stmt="fwd_bwd()", globals={"fwd_bwd": fwd_bwd}).timeit(ctx.bench_iters)

    if ctx.rank == 0:
        print(f"fwd: {fwd_t.mean*1000:.1f}ms | fwd+bwd: {fwd_bwd_t.mean*1000:.1f}ms", flush=True)


if __name__ == "__main__":
    ctx, args = setup()

    sp_llama = load_model_serialized(ctx, lambda m: make_sequence_parallel(m, ctx).to(ctx.device).eval())
    sp_llama = SPModelWrapper(
        sp_llama, gather_output=(ctx.world_size > 1) and (args.mode == "correctness")
    )  # to not allocate logits twice

    # 'eager' to still use ref model on cpu even if flash attn
    ref_llama = load_model_on_rank0(ctx, "eager")
    if ref_llama is not None:
        ref_llama.eval()

    if args.mode in ("correctness", "all"):
        test_correctness(sp_llama, ref_llama, DEFAULT_PROMPT, ctx)
    if args.mode in ("learning", "all"):
        test_learning(sp_llama, ctx)
    if args.mode in ("benchmark", "all"):
        benchmark(sp_llama, ctx)

    dist.destroy_process_group()