benchmark_low_bit_adam.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD 3-Clause license found in the
# LICENSE file in the root directory of this source tree.
# pip install timm wandb tqdm datasets bitsandbytes
#
# optional:
# - lpmm (4-bit optim): pip install yacs git+https://github.com/thu-ml/low-bit-optimizers.git
# - DeepSpeed (ZeRO-Offload):
#     sudo apt install libopenmpi-dev
#     LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu pip install mpi4py
#     DS_BUILD_CPU_ADAM=1 pip install deepspeed --no-cache-dir
#
# To fine-tune a pre-trained ViT-Base on resisc45 dataset with BF16 AMP, using default AdamW optimizer from PyTorch core
# python benchmark_low_bit_adam.py \
#   --model "timm/vit_base_patch16_224.augreg_in21k" \
#   --amp bf16 \
#   --optim AdamW
#
# See OPTIM_MAP for the available optimizer options
# To profile and export chrome trace, set --profile
# To enable cosine learning rate scheduler, set --cosine_lr_scheduler

import argparse
import json
import math
import time
from contextlib import nullcontext
from functools import partial

import bitsandbytes as bnb
import datasets
import timm
import torch
import torch.nn.functional as F
import wandb
from torch.utils.data import DataLoader
from torchvision.transforms import v2
from tqdm import tqdm

from torchao import optim
from torchao.utils import get_available_devices

_DEVICE = get_available_devices()[-1]
assert _DEVICE in ["cuda", "xpu"], "Benchmark currently only supports CUDA & XPU(BF16)"

OPTIM_MAP = dict(
    AdamW=partial(torch.optim.AdamW, fused=True),
    AdamW8bitBnb=bnb.optim.AdamW8bit,
    AdamW8bitAo=optim.AdamW8bit,
    AdamWFp8Ao=optim.AdamWFp8,
    AdamW4bitAo=optim.AdamW4bit,
)

try:
    import lpmm

    OPTIM_MAP.update(
        AdamW4bitLpmm=partial(lpmm.optim.AdamW, fused=True),
        AdamW4bitRank1Lpmm=partial(
            lpmm.optim.AdamW, qconfig=argparse.Namespace(scale_type="rank1")
        ),
    )

except ImportError:
    pass


class CosineSchedule:
    def __init__(self, lr: float, total_steps: int, warmup: float = 0.05) -> None:
        self.lr = lr
        self.final_lr = 0
        self.total_steps = total_steps
        self.warmup_steps = round(total_steps * warmup)

    def get_lr(self, step: int) -> float:
        if step < self.warmup_steps:
            return self.lr * step / self.warmup_steps
        if step < self.total_steps:
            progress = (step - self.warmup_steps) / (
                self.total_steps - self.warmup_steps
            )
            return self.final_lr + 0.5 * (self.lr - self.final_lr) * (
                1 + math.cos(progress * math.pi)
            )
        return self.final_lr


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--model_kwargs", type=json.loads, default=dict())
    parser.add_argument("--checkpoint_activations", action="store_true")

    parser.add_argument("--amp", default="none")
    parser.add_argument("--full_bf16", action="store_true")
    parser.add_argument("--channels_last", action="store_true")
    parser.add_argument("--compile", action="store_true")

    parser.add_argument("--n_epochs", type=int, default=10)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--n_workers", type=int, default=4)

    parser.add_argument("--optim", default="AdamW", choices=OPTIM_MAP.keys())
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight_decay", type=float, default=0)
    parser.add_argument("--optim_kwargs", type=json.loads, default=dict())
    parser.add_argument("--cosine_lr_scheduler", action="store_true")
    parser.add_argument(
        "--optim_cpu_offload", choices=["ao", "ao_offload_grads", "deepspeed"]
    )

    parser.add_argument("--project")
    parser.add_argument("--run_name", default="debug")
    parser.add_argument("--profile", action="store_true")
    parser.add_argument("--seed", type=int)
    return parser


def get_dloader(args, training: bool):
    transforms = [v2.ToImage()]

    if training:
        transforms.extend([v2.RandomResizedCrop(224), v2.RandomHorizontalFlip()])
    else:
        transforms.extend([v2.Resize(256), v2.CenterCrop(224)])

    transforms.append(v2.ToDtype(torch.float32, scale=True))
    transforms.append(
        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    )
    transforms = v2.Compose(transforms)

    # use dataset from HF so download is fast
    ds = datasets.load_dataset(
        "timm/resisc45", split="train" if training else "validation"
    )
    ds = ds.select_columns(["image", "label"])
    ds.set_transform(lambda x: dict(image=transforms(x["image"]), label=x["label"]))

    return DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=training,
        num_workers=args.n_workers,
        pin_memory=training,
        drop_last=training,
    )


def get_amp_ctx(amp, device):
    dtype = dict(bf16=torch.bfloat16, fp16=torch.float16, none=None)[amp]
    return torch.autocast(device, dtype=dtype, enabled=amp != "none")


@torch.no_grad()
def evaluate_model(model, args):
    model.eval()
    val_dloader = get_dloader(args, False)

    all_labels = []
    all_preds = []

    for batch in tqdm(val_dloader, dynamic_ncols=True, desc="Evaluating"):
        all_labels.append(batch["label"].clone())
        if args.full_bf16:
            batch["image"] = batch["image"].bfloat16()
        if args.channels_last:
            batch["image"] = batch["image"].to(memory_format=torch.channels_last)

        with get_amp_ctx(args.amp, _DEVICE):
            all_preds.append(model(batch["image"].to(_DEVICE)).argmax(1).cpu())

    all_labels = torch.cat(all_labels, dim=0)
    all_preds = torch.cat(all_preds, dim=0)

    acc = (all_labels == all_preds).float().mean()
    return acc


if __name__ == "__main__":
    args = get_parser().parse_args()

    if args.full_bf16:
        assert args.amp == "none", "When --full_bf16 is set, --amp must be none"
    if args.optim_cpu_offload == "deepspeed":
        assert (
            args.amp == "none"
        ), "When using DeepSpeed ZeRO-Offload, --amp must be none"
        assert (
            args.optim == "AdamW"
        ), "When using DeepSpeed ZeRO-Offload, --optim must be AdamW"
    if args.profile:
        args.n_epochs = 1
    if args.seed is not None:
        torch.manual_seed(args.seed)

    for k, v in vars(args).items():
        print(f"{k}: {v}")

    # wandb is only enabled when args.project is set and args.profile is False
    logger = wandb.init(
        project=args.project,
        name=args.run_name,
        config=args,
        dir="/tmp",
        mode="disabled" if args.project is None else None,
    )
    dloader = get_dloader(args, True)
    print(f"Train dataset: {len(dloader.dataset):,} images")

    model = timm.create_model(
        args.model, pretrained=True, num_classes=45, **args.model_kwargs
    )
    if args.checkpoint_activations:
        model.set_grad_checkpointing()
    if args.full_bf16:
        model.bfloat16()
    if args.channels_last:
        model.to(memory_format=torch.channels_last)
    model.to(_DEVICE)  # move model to DEVICE after optionally convert it to BF16
    if args.compile:
        model.compile(fullgraph=True)
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

    if args.optim_cpu_offload == "deepspeed":
        import deepspeed

        model, optim, _, _ = deepspeed.initialize(
            model=model,
            model_parameters=model.parameters(),
            config=dict(
                train_batch_size=args.batch_size,
                optimizer=dict(
                    type="Adam",
                    params=dict(
                        lr=args.lr,
                        weight_decay=args.weight_decay,
                        fp32_optimizer_states=False,
                        **args.optim_kwargs,
                    ),
                ),
                bf16=dict(enabled=args.full_bf16),
                zero_optimization=dict(
                    stage=2,  # requires ZeRO-2 to enable overlap_comm
                    overlap_comm=True,  # interleave grad D2H with backward
                    offload_optimizer=dict(device="cpu", pin_memory=True),
                ),
            ),
        )

    else:
        optim_cls = OPTIM_MAP[args.optim]

        if args.optim_cpu_offload == "ao":
            optim_cls = partial(optim.CPUOffloadOptimizer, optimizer_class=optim_cls)
        elif args.optim_cpu_offload == "ao_offload_grads":
            optim_cls = partial(
                optim.CPUOffloadOptimizer,
                optimizer_class=optim_cls,
                offload_gradients=True,
            )

        optim = optim_cls(
            model.parameters(),
            lr=args.lr,
            weight_decay=args.weight_decay,
            **args.optim_kwargs,
        )

    lr_schedule = CosineSchedule(args.lr, len(dloader) * args.n_epochs)
    grad_scaler = torch.amp.GradScaler(_DEVICE, enabled=args.amp == "fp16")
    log_interval = 10
    t0 = time.perf_counter()

    step = 0
    for epoch_idx in range(args.n_epochs):
        model.train()
        pbar = tqdm(
            dloader, dynamic_ncols=True, desc=f"Epoch {epoch_idx + 1}/{args.n_epochs}"
        )

        with torch.profiler.profile() if args.profile else nullcontext() as prof:
            for batch in pbar:
                if args.full_bf16:
                    batch["image"] = batch["image"].bfloat16()
                if args.channels_last:
                    batch["image"] = batch["image"].to(
                        memory_format=torch.channels_last
                    )

                with get_amp_ctx(args.amp, _DEVICE):
                    loss = F.cross_entropy(
                        model(batch["image"].to(_DEVICE)), batch["label"].to(_DEVICE)
                    )

                if args.optim_cpu_offload == "deepspeed":
                    model.backward(loss)
                else:
                    grad_scaler.scale(loss).backward()

                if args.cosine_lr_scheduler:
                    lr = lr_schedule.get_lr(step)
                    for param_group in optim.param_groups:
                        if isinstance(param_group["lr"], torch.Tensor):
                            param_group["lr"].fill_(lr)
                        else:
                            param_group["lr"] = lr

                if step % log_interval == 0:
                    log_dict = dict(loss=loss.item(), lr=optim.param_groups[0]["lr"])
                    if step > 0:
                        t1 = time.perf_counter()
                        log_dict["imgs_per_second"] = (
                            args.batch_size * log_interval / (t1 - t0)
                        )
                        t0 = t1
                    logger.log(log_dict, step=step)

                if args.optim_cpu_offload == "deepspeed":
                    model.step()
                else:
                    grad_scaler.step(optim)
                    grad_scaler.update()
                    optim.zero_grad()

                step += 1

                if args.profile and step == 5:
                    break

        if args.profile:
            prof.export_chrome_trace("trace.json")

        else:
            val_acc = evaluate_model(model, args)
            print(
                f"Epoch {epoch_idx + 1}/{args.n_epochs}: val_acc={val_acc.item() * 100:.2f}"
            )
            logger.log(dict(val_acc=val_acc), step=step)

    peak_mem = getattr(torch, _DEVICE).max_memory_allocated() / 1e9
    print(f"Max memory used: {peak_mem:.02f} GB")
    logger.log(dict(max_memory_allocated=peak_mem))