(
name: str,
batch_size: int = 8,
seq_len: int = 2048,
varlen: bool = False,
warmup_steps: int = 16,
steps: int = 32,
total_steps: int = 1024,
lr: float = 3e-4,
betas: Tuple[float] = (0.9, 0.95),
weight_decay: float = 0.1,
dtype: Optional[torch.dtype] = torch.bfloat16,
mixed_precision: str = 'bf16'
)
| 50 | |
| 51 | |
| 52 | def profile( |
| 53 | name: str, |
| 54 | batch_size: int = 8, |
| 55 | seq_len: int = 2048, |
| 56 | varlen: bool = False, |
| 57 | warmup_steps: int = 16, |
| 58 | steps: int = 32, |
| 59 | total_steps: int = 1024, |
| 60 | lr: float = 3e-4, |
| 61 | betas: Tuple[float] = (0.9, 0.95), |
| 62 | weight_decay: float = 0.1, |
| 63 | dtype: Optional[torch.dtype] = torch.bfloat16, |
| 64 | mixed_precision: str = 'bf16' |
| 65 | ): |
| 66 | device = torch.device('cuda') |
| 67 | config = configs[name] if name in configs else AutoConfig.from_pretrained(name) |
| 68 | model = AutoModelForCausalLM.from_config(config).cuda().to(dtype) |
| 69 | num_parameters = model.num_parameters() |
| 70 | print(f"Initializing {name} model from the config:\n{config}\n{model}") |
| 71 | print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})") |
| 72 | print(f"Allocated memory after initialization: {sizeof_fmt(memory_allocated(device))}") |
| 73 | |
| 74 | accelerator = Accelerator(mixed_precision=mixed_precision) |
| 75 | optimizer = AdamW( |
| 76 | model.parameters(), |
| 77 | lr=lr, |
| 78 | betas=betas, |
| 79 | weight_decay=weight_decay, |
| 80 | fused=True |
| 81 | ) |
| 82 | scheduler = get_cosine_schedule_with_warmup(optimizer, 0, total_steps) |
| 83 | |
| 84 | bar = trange(warmup_steps) |
| 85 | |
| 86 | model, optimizer, scheduler = accelerator.prepare(model, optimizer, scheduler) |
| 87 | torch.cuda.synchronize(device) |
| 88 | for _ in bar: |
| 89 | # forward pass |
| 90 | tokens, offsets = prepare_inputs( |
| 91 | batch_size=batch_size, |
| 92 | seq_len=seq_len, |
| 93 | varlen=varlen, |
| 94 | vocab_size=config.vocab_size, |
| 95 | device=device |
| 96 | ) |
| 97 | outputs = model(tokens, labels=tokens, offsets=offsets) |
| 98 | # backward pass |
| 99 | accelerator.backward(outputs.loss) |
| 100 | optimizer.step() |
| 101 | scheduler.step() |
| 102 | optimizer.zero_grad() |
| 103 | bar.set_description_str(f"Max memory allocated: {sizeof_fmt(max_memory_allocated(device))}") |
| 104 | |
| 105 | start, total_tokens = time.time(), 0 |
| 106 | bar = trange(steps) |
| 107 | torch.cuda.synchronize(device) |
| 108 | for _ in bar: |
| 109 | # forward pass |
no test coverage detected