utils.py

import os
import argparse
import torch
from dataclasses import dataclass
import torch.distributed as dist
import transformers
from torch.distributed.device_mesh import init_device_mesh
from transformers.models.llama.modeling_llama import LlamaConfig
from config import DEFAULT_MODEL, DEFAULT_ATTN_IMPL, DEFAULT_PROMPT


@dataclass
class Context:
    rank: int
    world_size: int
    device: torch.device
    dtype: torch.dtype
    config: LlamaConfig
    tokenizer: transformers.PreTrainedTokenizerBase
    model_name: str
    attn_impl: str
    length: int
    device_mesh: torch.distributed.device_mesh.DeviceMesh
    bench_iters: int

    @property
    def is_first(self):
        return self.rank == 0

    @property
    def is_last(self):
        return self.rank == self.world_size - 1


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--length", type=int, default=2048)
    parser.add_argument("--mode", choices=["benchmark", "correctness", "learning", "all"], default="all")
    parser.add_argument("--dtype", choices=["float32", "bfloat16"], default="bfloat16")
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--dtensor", action="store_true", default=False)
    parser.add_argument("--bench-iters", type=int, default=7)
    return parser.parse_args()


def load_config(model_name, attn_impl=DEFAULT_ATTN_IMPL):
    config = LlamaConfig.from_pretrained(model_name, attn_implementation=attn_impl)
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
    return config, tokenizer


def init_distributed():
    rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    device = torch.device(f"cuda:{rank}")
    dist.init_process_group("nccl", device_id=device)
    torch.manual_seed(1337)
    device_mesh = init_device_mesh(device_type="cuda", mesh_shape=(world_size,))
    return rank, world_size, device, device_mesh


def setup():
    args = parse_args()
    rank, world_size, device, device_mesh = init_distributed()
    dtype = getattr(torch, args.dtype)
    model_name = args.model or DEFAULT_MODEL_NAME
    if rank == 0 and dtype == torch.float32:
        print("float32 is incompatible with flash attn, using 'eager'")
    attn_impl = "eager" if dtype == torch.float32 else DEFAULT_ATTN_IMPL
    config, tokenizer = load_config(model_name, attn_impl)
    ctx = Context(
        rank=rank,
        world_size=world_size,
        device=device,
        dtype=dtype,
        config=config,
        tokenizer=tokenizer,
        model_name=model_name,
        attn_impl=attn_impl,
        length=args.length,
        device_mesh=device_mesh,
        bench_iters=args.bench_iters,
    )
    return ctx, args


def load_model_serialized(ctx, transform_fn=None):
    for r in range(ctx.world_size):
        dist.barrier()
        if r != ctx.rank:
            continue
        model = transformers.LlamaForCausalLM.from_pretrained(
            ctx.model_name, attn_implementation=ctx.attn_impl, torch_dtype=ctx.dtype
        )
        if transform_fn:
            model = transform_fn(model)
    return model


def load_model_on_rank0(ctx, attn_impl=None):
    if attn_impl is None:
        attn_impl = ctx.attn_impl
    if ctx.rank != 0:
        return None
    ref = transformers.LlamaForCausalLM.from_pretrained(
        ctx.model_name, attn_implementation=attn_impl, dtype=ctx.dtype
    )
    return ref