...
  File "/home/yyc/.local/miniconda3/envs/zero_env/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 611, in run
    return model(new_inputs)
  File "/home/yyc/.local/miniconda3/envs/zero_env/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 892, in _run_from_cache
    return compiled_graph.compiled_artifact(inputs)
  File "/tmp/torchinductor_yyc/u6/cu6jjzq2vnxkdniqmnw3r7htiuocckvp7ou6qiiyn2ofrtdy7rz7.py", line 503, in call
    buf12 = aten._scaled_dot_product_efficient_attention(buf7, buf8, buf9, buf11, True)
  File "/home/yyc/.local/miniconda3/envs/zero_env/lib/python3.10/site-packages/torch/_ops.py", line 755, in __call__
    return self._op(*args, **(kwargs or {}))
RuntimeError: invalid dtype for bias - should match query's dtype

import argparse
import gc
import json
import os
import time

import numpy as np
import torch
import torch.distributed as dist
import transformers
from accelerate import init_empty_weights
from datasets import Dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy
from transformers import (AutoConfig, AutoModelForCausalLM, LlamaConfig,
                          PreTrainedModel, Trainer, TrainerCallback,
                          TrainingArguments)

assert transformers.__version__ >= '4.36', 'requires transformers 4.36+ to enable sdpa by default'

parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='gpt2')
parser.add_argument('--mbs', type=int, default=1)
parser.add_argument('--ga', type=int, default=1)
parser.add_argument('--seq-len', type=int, default=512)
parser.add_argument('--iters', type=int, default=100)
parser.add_argument('--cc', default=None)
parser.add_argument('--layers', type=int, default=2)
parser.add_argument('--embd', type=int, default=-1)
parser.add_argument('--vocab-size', type=int, default=-1)
parser.add_argument('--zero-stage', '--zero', type=int, default=0)
parser.add_argument('--overlap-comm', action="store_true", default=False)
parser.add_argument('--local-rank', type=int, default=-1)
parser.add_argument('--print-gc', action="store_true", default=False)



class ThrputCb(TrainerCallback):
    def __init__(self, iters, seq_len):
        self.t0 = time.perf_counter()
        self.total_iters = iters
        self.seq_len = seq_len
        self.iter = 0
        self.smooth_thrput = 0
        self.smooth_toks_per_dev = 0

    def on_log(self, args, state, control, logs, **kwargs):
        _ = logs.pop("total_flos", None)
        self.iter += 1
        if not state.is_local_process_zero:
            return
        if self.iter > self.total_iters:
            print(
                f'🎉 {self.smooth_thrput:.2f} samples/sec, {self.smooth_toks_per_dev:.2f} toks/GPU')
            return
        world_size = dist.get_world_size() if dist.is_initialized() else 1
        now = time.perf_counter()
        elapsed = now - self.t0
        self.t0 = now
        thrput = world_size * args.per_device_train_batch_size * \
            args.gradient_accumulation_steps / elapsed
        self.smooth_thrput = self.smooth_thrput * 0.9 + thrput * 0.1
        toks_per_dev = thrput / world_size * self.seq_len
        self.smooth_toks_per_dev = self.smooth_toks_per_dev * 0.9 + toks_per_dev * 0.1
        print(f'🐣 Iter {self.iter}/{self.total_iters}: {elapsed*1000:.2f}ms, {thrput:.2f}({self.smooth_thrput:.2f})samples/sec, {toks_per_dev:.2f}({self.smooth_toks_per_dev:.2f})toks/GPU, gc={gc.get_count()}', end=' ')


def print_rank0(s):
    if not dist.is_initialized() or dist.get_rank() == 0:
        print(s)


def load_model(model_id: str, layers, embd, vocab_size, zero) -> PreTrainedModel:
    # Reduce dependence on the network environment for llama
    if model_id == 'meta-llama/Llama-2-7b-hf':
        config = LlamaConfig(hidden_size=4096, num_hidden_layers=32,
                             num_attention_heads=32, num_key_value_heads=32, intermediate_size=11008)
    elif model_id == 'meta-llama/Llama-2-13b-hf':
        config = LlamaConfig(hidden_size=5120, num_hidden_layers=40,
                             num_attention_heads=40, num_key_value_heads=40, intermediate_size=13824)
    elif model_id == 'meta-llama/Llama-2-30b-hf':  # not official
        config = LlamaConfig(hidden_size=6400, num_hidden_layers=60,
                             num_attention_heads=40, num_key_value_heads=40, intermediate_size=17280)
    elif model_id == 'meta-llama/Llama-2-70b-hf':
        config = LlamaConfig(hidden_size=8192, num_hidden_layers=80,
                             num_attention_heads=64, num_key_value_heads=8, intermediate_size=28672)
    else:
        if os.path.exists(f'./models/{model_id}'):
            model_id = f'./models/{model_id}'
            print_rank0(f'💾 Loading model from local {model_id}')
        config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)

    if layers != -1: 
        if hasattr(config, 'num_hidden_layers'):
            config.num_hidden_layers = layers
        elif hasattr(config, 'num_layers'):
            config.num_layers = layers
    if embd != -1:
        if hasattr(config, 'hidden_size'):
            config.hidden_size = embd
        elif hasattr(config, 'n_embd'):
            config.n_embd = embd
    if vocab_size != -1:
        if hasattr(config, 'vocab_size'):
            config.vocab_size = vocab_size

    if zero:
        import deepspeed
        with deepspeed.zero.Init():
            model = AutoModelForCausalLM.from_config(
                config, trust_remote_code=True)
    else:
        with init_empty_weights():
            model = AutoModelForCausalLM.from_config(
                config, trust_remote_code=True).to_empty(device='cpu')
    model.config.pad_token_id = model.config.eos_token_id
    return model


def guess_transformer_layer_class(model: torch.nn.Module):
    mlist = [m for m in model.modules() if isinstance(m, torch.nn.ModuleList)]
    return mlist[0][0].__class__


class CompileOptimizerTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(), lr=torch.scalar_tensor(1e-5), foreach=True, capturable=True)
        self.lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=1000, eta_min=1e-6)
        self.optimizer.step = torch.compile(self.optimizer.step)


def train(model, args):
    # deepspeed_dict = json.loads(DEEPSPEED_TEMPLATE)
    # deepspeed_dict['zero_optimization']['stage'] = args.zero_stage
    # deepspeed_dict['zero_optimization']['overlap_comm'] = args.overlap_comm
    # deepspeed_dict['wall_clock_breakdown'] = True

    input_ids = np.random.randint(100, 30000, (1000, args.seq_len))
    data_set = Dataset.from_dict({
        "input_ids": input_ids,
        "labels": input_ids
    })

    training_args = TrainingArguments(
        log_level='info',
        per_device_train_batch_size=args.mbs,
        gradient_accumulation_steps=args.ga,
        max_steps=args.iters,
        save_steps=10e9,
        logging_steps=1,
        output_dir='./tmp',
        disable_tqdm=True,
        bf16=True,
        torch_compile_backend=args.cc,
    )

    trainer_cls = CompileOptimizerTrainer if args.cc and args.zero_stage == 0 else Trainer
    trainer = trainer_cls(model, args=training_args,
                          train_dataset=data_set,
                          callbacks=[ThrputCb(args.iters, args.seq_len)])
    trainer.train()


if __name__ == '__main__':
    if 'RANK' in os.environ and not dist.is_initialized():
        dist.init_process_group(backend='nccl', init_method='env://')
    args = parser.parse_args()
    print_rank0(f'args: {args}')
    if args.zero_stage != 0 and args.fsdp:
        raise ValueError('FSDP and ZeRO are mutually exclusive')
    if args.print_gc:
        gc.set_debug(gc.DEBUG_STATS)
    model: PreTrainedModel = load_model(
        args.model, args.layers, args.embd, args.vocab_size, args.zero_stage != 0)
    print_rank0(f'{model}\nparameters: {model.num_parameters()/1e9:.3f} B')
    train(model, args)