tp.py

import copy

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

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

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


class ComputeWithAllReduce(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tp_shard: nn.Module, input: torch.Tensor, args, kwargs):
        input = input.detach().requires_grad_(input.requires_grad)
        ctx.save_for_backward(input)
        ctx._tp_shard = tp_shard
        ctx._kwargs = kwargs
        ctx._args = args
        output = tp_shard(input, *args, **kwargs)
        dist.all_reduce(output[0] if isinstance(output, tuple) else output)
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor, *rest_grads):
        with torch.enable_grad():
            output = ctx._tp_shard(ctx.saved_tensors[0], *ctx._args, **ctx._kwargs)
            out0 = output[0] if isinstance(output, tuple) else output
            out0.backward(grad_output)
        dist.all_reduce(ctx.saved_tensors[0].grad)
        return None, ctx.saved_tensors[0].grad, None, None


class AllReduceModule(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 ComputeWithAllReduce.apply(self.module, input, rest_args, kwargs)


def make_tensor_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, AllReduceModule(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, AllReduceModule(tp_mlp))
        else:
            make_tensor_parallel(child, ctx)

    return model


def make_tensor_parallel_dtensor(model, ctx):
    """Alternative implementation with the help of DTensor."""
    l = list(model.named_children())
    for name, child in l:
        if isinstance(child, LlamaAttention):
            tp_attn = tp.parallelize_module(
                child,
                ctx.device_mesh,
                parallelize_plan={
                    "q_proj": tp.ColwiseParallel(),
                    "k_proj": tp.ColwiseParallel(),
                    "v_proj": tp.ColwiseParallel(),
                    "o_proj": tp.RowwiseParallel(),
                },
            )
            setattr(model, name, tp_attn)
        elif isinstance(child, LlamaMLP):
            tp_mlp = tp.parallelize_module(
                child,
                ctx.device_mesh,
                parallelize_plan={
                    "gate_proj": tp.ColwiseParallel(),
                    "up_proj": tp.ColwiseParallel(),
                    "down_proj": tp.RowwiseParallel(),
                },
            )
            setattr(model, name, tp_mlp)
        else:
            make_tensor_parallel_dtensor(child, ctx)
    return model


def test_correctness(tp_llama, ref_llama, prompt, ctx):
    input_ids = ctx.tokenizer(prompt, return_tensors="pt")["input_ids"]
    tp_embeds = tp_llama.model.embed_tokens(input_ids.to(ctx.device)).detach().requires_grad_(True)
    tp_logits = tp_llama(inputs_embeds=tp_embeds, use_cache=False).logits
    tp_logits.mean().backward()

    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 = (tp_logits.cpu() - ref_logits).detach().abs().max()
        grad_diff = (tp_embeds.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(tp_model, ctx):
    for p in tp_model.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 = tp_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):
        loss = tp_model(inputs_embeds=embeds, labels=labels, use_cache=False).loss
        opt.zero_grad()
        loss.backward()
        opt.step()

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


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

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

    def fwd_bwd():
        tp_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()

    if args.dtensor:
        model = load_model_serialized(ctx, lambda m: m.to(ctx.device).eval())
        tp_llama = make_tensor_parallel_dtensor(model, ctx)  # it is sync op
        # we need checkpointing for fair comparison
        for layer in tp_llama.model.layers:
            for name in ("self_attn", "mlp"):
                module = getattr(layer, name)
                original_forward = module.forward
                module.forward = lambda *args, _fn=original_forward, **kwargs: torch.utils.checkpoint.checkpoint(
                    _fn, *args, use_reentrant=False, **kwargs
                )
    else:
        tp_llama = load_model_serialized(ctx, lambda m: make_tensor_parallel(m, ctx).to(ctx.device).eval())

    # '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(tp_llama, ref_llama, DEFAULT_PROMPT, ctx)
    if args.mode in ("learning", "all"):
        test_learning(tp_llama, ctx)
    if args.mode in ("benchmark", "all"):
        benchmark(tp_llama, ctx)

    dist.destroy_process_group()