main_utils.py

import os
import glob
import json
import tqdm
import math
import numpy as np
from torch.distributed.fsdp import FullStateDictConfig
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import StateDictType
# from lightning.fabric.strategies import FSDPStrategy # REMOVE ME

from model_utils.modeling_llama import LlamaForCausalLM


def load_jsonl_examples(filename,
                        n_examples,
                        shuffle,
                        global_micro_batch_size,
                        global_rank,
                        world_size):
    example_idxes = np.random.permutation(n_examples) if shuffle \
        else np.arange(n_examples)

    n_examples = n_examples // global_micro_batch_size * global_micro_batch_size
    example_idxes = example_idxes[global_rank:n_examples:world_size]

    examples = {idx: None for idx in example_idxes}
    for example_idx, line in tqdm.tqdm(
            enumerate(open(filename)), desc=f'loading {filename}'):
        if example_idx in examples:
            examples[example_idx] = json.loads(line)

    return [examples[idx] for idx in example_idxes]


def get_cosine_lr_decay_fn(total_steps,
                           warmup_steps,
                           learning_rate,
                           end_learning_rate):
    def cosine_with_warmup_lr(step):
        if step < warmup_steps:
            return learning_rate * step / warmup_steps
        elif step > total_steps:
            return end_learning_rate

        decay_ratio = (step - warmup_steps) / (total_steps - warmup_steps)
        assert 0 <= decay_ratio <= 1
        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
        return end_learning_rate + coeff * (learning_rate - end_learning_rate)

    return cosine_with_warmup_lr


def get_grad_norm(model):
    square_sum = 0.
    for param in model.parameters():
        if param.grad is not None:
            square_sum += param.grad.detach().data.norm(2).item() ** 2
    return square_sum ** 0.5


# def save_checkpoint(fabric, tokenizer, model, optimizer, save_dir):
#     assert isinstance(fabric.strategy, FSDPStrategy)

#     save_policy = FullStateDictConfig(
#         offload_to_cpu=(fabric.world_size > 1), rank0_only=True)
#     with FSDP.state_dict_type(
#             model,
#             state_dict_type=StateDictType.FULL_STATE_DICT,
#             state_dict_config=save_policy):
#         state_dict = model._forward_module.state_dict()
    
#     if fabric.global_rank == 0:
#         tokenizer.save_pretrained(save_dir)
#         assert isinstance(model.module, LlamaForCausalLM)
#         model.module.save_pretrained(
#             save_dir, state_dict=state_dict, safe_serialization=False)

#     fabric.barrier()
#     fabric.save(
#         path=f'{save_dir}/fabric_ckpt',
#         state={'model': model, 'optimizer': optimizer})

def save_checkpoint(fabric, tokenizer, model, optimizer, save_dir): # REMOVE ME
    if fabric.global_rank == 0:
        tokenizer.save_pretrained(save_dir)
        assert isinstance(model.module, LlamaForCausalLM)
        model.module.save_pretrained(save_dir, safe_serialization=False)

    fabric.barrier()
    try:
        fabric.save(
            path=f'{save_dir}/fabric_ckpt',
            state={'model': model.state_dict(), 'optimizer': optimizer.state_dict()})
    except OSError as e:
        print(f"Failed to save checkpoint to {save_dir}/fabric_ckpt: {e}")
        
def get_last_ckpt_idx(workdir):
    last_ckpt_idx = -1
    for ckpt_dir in glob.glob(f'{workdir}/ckpt_*'):
        ckpt_idx = int(ckpt_dir.split('_')[-1])
        if ckpt_idx > last_ckpt_idx:
            last_ckpt_idx = ckpt_idx

    return last_ckpt_idx