src/lgmodeling/finetune.py

import os
import math
import time
import copy
import json
import torch
import jax
import argparse
import optax
import wandb
import shutil
import itertools
from tqdm import tqdm
from typing import Any, Dict, List
from copy import deepcopy
from datasets import Dataset
from datetime import timedelta
from data_utils import get_lm_corpus
from flax.jax_utils import replicate, unreplicate
from flax.core.frozen_dict import freeze, unfreeze
from flax.training import train_state, checkpoints
from flax.traverse_util import flatten_dict, unflatten_dict
from flax.training.common_utils import get_metrics, onehot, shard
from transformers.models.gpt2.modeling_flax_gpt2 import GPT2Config
from lmc_model import LMCFlaxGPT2LMHeadModel, print_model
from data_utils import get_lm_corpus
import jax.numpy as jnp
os.environ["WANDB_API_KEY"] = "fc72050bcc0dc7f7502b5416938f8bd0c4b30fc7"
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
os.environ["JAX_DEFAULT_MATMUL_PRECISION"] = "highest" 
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
def remove_old_dirs_with_prefix(save_path, prefix, keep_step):
    for fname in os.listdir(save_path):
        if fname.startswith(prefix) and not fname.endswith(str(keep_step)):
            full_path = os.path.join(save_path, fname)
            if os.path.isdir(full_path):
                shutil.rmtree(full_path)



def prepare_lm_batch(data: torch.Tensor, target: torch.Tensor) -> Dict[str, Any]:
    """
    Convert and shard a language modeling batch from PyTorch to JAX.
    Args:
        data (torch.Tensor): Input data of shape (seq_len, batch)
        target (torch.Tensor): Target data of shape (seq_len, batch)
    Returns:
        Dict[str, jnp.ndarray]: Dict with 'data' and 'target', both sharded
            with shape (n_devices, batch_per_device, seq_len)
    """
    # Transpose to (batch, seq_len), then convert to jnp arrays
    input_ids =  jnp.array(data.T)
    target = jnp.array(target.T)
    # Shard across devices
    return {'input_ids': shard(input_ids),'target': shard(target)}
def decay_mask_fn(params):
    flat_params = flatten_dict(params)
    flat_mask = {
        path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")])
        for path in flat_params
    }
    return unflatten_dict(flat_mask)
# ---------- Training Utilities ----------
def get_trainable_mask(params, config):
    def is_trainable_param(keys):
        # Match MoE parameters in transformer/h/{moe_idx}/mlp/(gate|routed_experts_*)
        if len(keys) < 4: return False
        if keys[0] == "transformer" and keys[1] == "h" and int(keys[2]) in config.lmc_layer_indices and keys[3] == "attn": return True
        if config.finetune_mlp == True and keys[0] == "transformer" and keys[1] == "h" and int(keys[2]) in config.lmc_layer_indices and keys[3] == "moe": return True
        return False
    def label_fn(path, _):
        keys = [str(k.key) for k in path]
        return "trainable" if is_trainable_param(keys) else "frozen"
    return jax.tree_util.tree_map_with_path(label_fn, params)
def pretrained2finetune_params(pretrained_params, finetune_params, config):
    pretrained_params = unfreeze(pretrained_params)
    finetune_params = unfreeze(finetune_params)
    # 1. Copy top-level embeddings and final layer norm
    finetune_params["transformer"]["wte"] = copy.deepcopy(pretrained_params["transformer"]["wte"])
    if config.position_embeddings == "learnable":
        finetune_params["transformer"]["wpe"] = copy.deepcopy(pretrained_params["transformer"]["wpe"])
    finetune_params["transformer"]["ln_f"] = copy.deepcopy(pretrained_params["transformer"]["ln_f"])
    # 2. Copy encoder layers
    for i in range(config.num_hidden_layers):
        str_i = str(i)
        if i in config.lmc_layer_indices:
            # Handle MoE layer: copy attention and norms from pretrained
            ref_layer = copy.deepcopy(pretrained_params["transformer"]["h"][str_i])
            target_layer = copy.deepcopy(finetune_params["transformer"]["h"][str_i])
            target_layer["ln_1"] = copy.deepcopy(ref_layer["ln_1"])
            # target_layer["attn"] = copy.deepcopy(ref_layer["attn"])
            target_layer["ln_2"] = copy.deepcopy(ref_layer["ln_2"])
            if config.finetune_mlp == False:
                target_layer["moe"] = copy.deepcopy(ref_layer["moe"])
        else:
            # Standard block, copy all directly
            finetune_params["transformer"]["h"][str_i] = copy.deepcopy(pretrained_params["transformer"]["h"][str_i])
    return freeze(finetune_params)

def main(args: argparse.Namespace):
    if os.path.exists(args.model_path):
        config = GPT2Config.from_pretrained(os.path.dirname(args.model_path))
        pretrained_model = LMCFlaxGPT2LMHeadModel(config)
    else:
        raise FileNotFoundError(f"Config directory does not exist: {os.path.dirname(args.model_path)}")
    os.makedirs(args.wandb_run_dir, exist_ok=True)
    wandb.init(
        project=args.wandb_project,
        entity=args.wandb_entity,
        group=args.wandb_group,
        name=f"finetune-{config.position_embeddings}-indice{','.join(str(i) for i in args.lmc_layer_indices)}-heads{args.n_head}"
            f"-shared{config.num_shared_experts}-routed{config.num_routed_experts}-topk{config.topk}-mlp{str(args.finetune_mlp)}-seed{args.seed}",
        save_code=True
    )
    save_path = os.path.join(args.model_save_dir,wandb.run.name)
    wandb.config = dict(vars(args))
    corpus = get_lm_corpus(args.data_path, args.dataset)
    ntokens = len(corpus.vocab)
    args.n_token = ntokens
    eval_batch_size = 12
    tr_iter = corpus.get_iterator('train', args.batch_size, args.tgt_len, ext_len=args.ext_len)
    va_iter = corpus.get_iterator('valid', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len)
    te_iter = corpus.get_iterator('test', eval_batch_size, args.eval_tgt_len, ext_len=args.ext_len)
    lmc_config = copy.deepcopy(config)
    lmc_config.n_head = args.n_head
    lmc_config.routed_scaling_factor = args.routed_scaling_factor
    config.lmc_config = lmc_config
    config.lmc_layer_indices = args.lmc_layer_indices
    config.finetune_mlp = args.finetune_mlp
    pretrained_params = checkpoints.restore_checkpoint(ckpt_dir=args.model_path, target={"params": pretrained_model.params})["params"]
    pretrained_model.params = pretrained_params
    # --- Initialize fine-tuning model ---
    model = LMCFlaxGPT2LMHeadModel(config,input_shape=(1, args.tgt_len),seed=args.seed,dtype=jnp.dtype(args.dtype),)
    model.config.save_pretrained(save_path)
    print_model(model.params)
    model.params = pretrained2finetune_params(pretrained_model.params,model.params,config)
    # model = pretrained_model
    label_mask = get_trainable_mask(model.params,config)
    print(json.dumps(label_mask, indent=2))
    num_train_steps = args.max_step
    lr_schedule =optax.warmup_cosine_decay_schedule(
        init_value=0.0,
        peak_value=args.learning_rate,
        warmup_steps=args.warmup_step,
        decay_steps=args.max_step,
        end_value=args.eta_min,
    )    
    tx = optax.multi_transform(
        transforms={
            'trainable': optax.adamw(
                learning_rate=lr_schedule,
                b1=args.adamw_beta1,
                b2=args.adamw_beta2,
                eps=args.adamw_eps,
                weight_decay=args.weight_decay_rate
            ),
            'frozen': optax.set_to_zero()
        },
        param_labels=label_mask  
    )
    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx)
    latest_train_step = state.step
    def train_step(state, batch, dropout_rng):
        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
        def loss_fn(params):
            labels = batch.pop("target")
            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
            loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean()
            return loss
        grad_fn = jax.value_and_grad(loss_fn)
        loss, grads = grad_fn(state.params)
        grads = jax.lax.pmean(grads, axis_name="batch")
        new_state = state.apply_gradients(grads=grads)
        metrics = {"loss": loss,"learning_rate": lr_schedule(state.step)}
        metrics = jax.lax.pmean(metrics, axis_name="batch")
        return new_state, metrics, new_dropout_rng
    def eval_step(state, batch):
        labels = batch.pop("target")
        logits = model(**batch, params=state.params, train=False)[0]
        loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean()
        metrics = {"eval_loss": loss}
        metrics = jax.lax.pmean(metrics, axis_name="batch")
        return metrics
    parallel_train_step = jax.pmap(train_step, "batch")
    parallel_eval_step = jax.pmap(eval_step, "batch")
    state = replicate(state)
    rng = jax.random.PRNGKey(args.seed)
    train_metrics_stack = []
    train_step = int(jax.device_get(unreplicate(state.step)))
    train_loss = 0.0
    best_val_loss = float("inf")
    log_start_time = time.time()
    eval_start_time = time.time()
    # ###JUST FOR TESTING####
    eval_results = []
    for eval_data, eval_target, _ in va_iter:
        eval_batch = prepare_lm_batch(eval_data, eval_target)
        eval_metric = parallel_eval_step(state, eval_batch)
        eval_results.append(eval_metric)
    eval_metrics = get_metrics(eval_results)
    eval_metrics = unreplicate(eval_metrics)
    eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics)
    val_loss = float(eval_metrics["eval_loss"])
    val_ppl = math.exp(val_loss)
    print("-" * 100)
    print(
        f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | "
        f"time: {time.time() - eval_start_time:5.2f}s | "
        f"valid loss {val_loss:5.2f} | valid ppl {val_ppl:9.3f}"
    )
    print("-" * 100)
    # #### START FINETUNING ####
    print("Starting training...")
    print(f"JAX devices: {jax.devices()}")
    print(f"Using {jax.local_device_count()} devices")
    for epoch in itertools.count(start=1):
        print(f"Epoch {epoch}")
        dropout_rngs = jax.random.split(rng, jax.local_device_count())
        train_iter = tr_iter.get_varlen_iter() if getattr(args, "varlen", False) else tr_iter
        train_metrics_stack = []
        for batch_idx, (data, target, seq_len) in enumerate(tqdm(train_iter)):
            if train_step >= args.max_step:
                break
            # Prepare and shard batch
            batch = prepare_lm_batch(data, target)
            # Run train step
            state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
            train_metrics_stack.append(train_metric)
            train_step += 1
            # Logging
            if train_step % args.logging_frequency == 0:
                train_metrics = get_metrics(train_metrics_stack)
                train_metrics = unreplicate(train_metrics)
                train_metrics = jax.tree_util.tree_map(lambda x: x.mean(), train_metrics)
                train_metrics_stack = []
                loss = float(train_metrics["loss"])
                ppl = math.exp(loss)
                bpc = loss/math.log(2)
                curr_lr = float(lr_schedule(train_step))
                elapsed = time.time() - log_start_time
                if(args.dataset in ["wt103","lm1b"]):
                    print(
                        f"| epoch {epoch:3d} step {train_step:8d} | "
                        f"{batch_idx+1:6d} batches | lr {curr_lr:.3g} "
                        f"| ms/batch {elapsed * 1000 / args.logging_frequency:5.2f} | "
                        f"loss {loss:5.2f} | ppl {ppl:9.3f}"
                    )
                    wandb.log({"loss": loss,"ppl": ppl,"learning_rate": curr_lr}, step=train_step)
                elif(args.dataset in ["enwik8","text8"]):
                    print(
                        f"| epoch {epoch:3d} step {train_step:8d} | "
                        f"{batch_idx+1:6d} batches | lr {curr_lr:.3g} "
                        f"| ms/batch {elapsed * 1000 / args.logging_frequency:5.2f} | "
                        f"loss {loss:5.2f} | bpc {bpc:9.3f}"
                    )
                    wandb.log({"loss": loss,"bpc": bpc,"learning_rate": curr_lr}, step=train_step)
                log_start_time = time.time()
            # Evaluation
            if train_step % args.eval_frequency == 0:
                eval_results = []
                for eval_data, eval_target, _ in va_iter:
                    eval_batch = prepare_lm_batch(eval_data, eval_target)
                    eval_metric = parallel_eval_step(state, eval_batch)
                    eval_results.append(eval_metric)
                eval_metrics = get_metrics(eval_results)
                eval_metrics = unreplicate(eval_metrics)
                eval_metrics = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics)
 
                val_loss = float(eval_metrics["eval_loss"])
                val_ppl = math.exp(val_loss)
                val_bpc = val_loss/math.log(2)
                print("-" * 100)
                if(args.dataset in ["wt103","lm1b"]):
                    print(
                        f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | "
                        f"time: {time.time() - eval_start_time:5.2f}s | "
                        f"valid loss {val_loss:5.2f} | valid ppl {val_ppl:9.3f}"
                    )
                    wandb.log({"eval_loss": val_loss,"eval_ppl": val_ppl}, step=train_step)
                elif(args.dataset in ["enwik8","text8"]):
                    print(
                        f"| Eval {train_step // args.eval_frequency:3d} at step {train_step:8d} | "
                        f"time: {time.time() - eval_start_time:5.2f}s | "
                        f"valid loss {val_loss:5.2f} | valid bpc {val_bpc:9.3f}"
                    )
                    wandb.log({"eval_loss": val_loss,"eval_bpc": val_bpc}, step=train_step)
                print("-" * 100)
                # Save best checkpoint
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    model.params = unreplicate(state).params
                    best_dir = os.path.join(save_path, f"best_{train_step}")
                    model.save_pretrained(best_dir)
                    print(f"✅ Best model saved at step {train_step}")
                    remove_old_dirs_with_prefix(save_path, "best_", train_step)
                eval_start_time = time.time()
            # Periodic checkpoint
            if train_step % args.save_frequency == 0:
                model.params = unreplicate(state).params
                last_dir = os.path.join(save_path, f"last_{train_step}")
                model.save_pretrained(last_dir)
                print(f"💾 Checkpoint saved at step {train_step}")
                remove_old_dirs_with_prefix(save_path, "last_", train_step)
                # checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=int(jax.device_get(unreplicate(state.step))),prefix="last_",keep=1,overwrite=True)                    
                print(f"Checkpoint saved at step {save_path}")
        if train_step >= args.max_step:
            print("-" * 100)
            print("End of training")
            break


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, default="", help="Path of Pretrained Model")
    parser.add_argument("--position-embeddings", type=str, default='sinusoidal')
    parser.add_argument("--rotary-dim", type=int,default=0,help=' Number of dimensions in the embedding that Rotary Position Embedding is applied to..')
    parser.add_argument("--num-shared-experts", type=int, default = 1)
    parser.add_argument("--num-routed-experts", type=int, default = 0)
    parser.add_argument("--topk",  type=int, default = 0)
    parser.add_argument('--n_layer', type=int, default=12,help='number of total layers')
    parser.add_argument('--n_head', type=int, default=10,help='number of heads')
    parser.add_argument('--n_embd', type=int, default=500,help='model dimension')
    parser.add_argument('--n_inner', type=int, default=1000,help='inner dimension in FF')
    parser.add_argument('--q_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for query projections.')
    parser.add_argument('--qk_rope_head_dim', type=int, default=64,help='Head dimension used for RoPE on query/key.')
    parser.add_argument('--kv_lora_rank', type=int, default=8,help='Rank of the LoRA adaptation for key/value projections.')
    parser.add_argument('--v_head_dim', type=int, default=64,help='Head dimension used for value projections.')
    parser.add_argument('--qk_nope_head_dim', type=int, default=64,help='Head dimension for NOPE (non-position encoding) on query/key.')
    parser.add_argument("--attention-bias", action="store_true",help='Use Bias in Attention.')
    parser.add_argument("--finetune-mlp",action="store_true",help="Enable fine-tuning for the MLP. Default is False.")
    parser.add_argument('--routed-scaling-factor', type=float, default=1.0,help='')
    parser.add_argument("--lmc-layer-indices",type=int,nargs="*",default=[],help="List of lmc layer indices (optional, default: empty list)")
    parser.add_argument("--seed", type=int, default=0, help="random seed for RNG state")
    parser.add_argument("--data-path", type=str, default="/cm/archive/vinhbk1/datasets/wikitext103", help="train datset paths (multiple paths)")
    parser.add_argument('--dataset', type=str, default='wt103',choices=['wt103', 'lm1b', 'enwik8', 'text8'],help='dataset name')
    parser.add_argument('--max_step', type=int, default=500000,help='upper epoch limit')
    parser.add_argument('--warmup_step', type=int, default=2000,help='upper epoch limit')
    parser.add_argument("--batch-size", type=int, default=96, help="train, eval batch size (batch size will be devided by device count)")
    parser.add_argument('--tgt_len', type=int, default=256,help='number of tokens to predict')
    parser.add_argument('--eval_tgt_len', type=int, default=256,help='number of tokens to predict for evaluation')
    parser.add_argument('--ext_len', type=int, default=0,help='length of the extended context')
    parser.add_argument('--mem_len', type=int, default=0,help='length of the retained previous heads')
    parser.add_argument("--learning-rate", type=float, default=0.00025, help="learning rate")
    parser.add_argument("--weight-decay-rate", type=float, default=0.01, help="weight deacy rate for lr scheduler")
    parser.add_argument('--eta_min', type=float, default=1.0e-8,help='min learning rate for cosine scheduler')
    parser.add_argument("--adamw-beta1", type=float, default=0.9)
    parser.add_argument("--adamw-beta2", type=float, default=0.999)
    parser.add_argument("--adamw-eps", type=float, default=1e-8)
    parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype")
    parser.add_argument("--wandb-entity", default="", help="wandb entity for logging")
    parser.add_argument("--wandb-group", default="", help="wandb group for logging")
    parser.add_argument("--wandb-project", default="GPT2-Wikitext103", help="wandb project name for logging")
    parser.add_argument("--wandb-run-dir", default=".wandb", help="wandb run dir")
    parser.add_argument("--logging-frequency", type=int, default=200, help="do logging every logging_frequency step")
    parser.add_argument("--eval-frequency", type=int, default=4000, help="do evalution every eval_frequency step")
    parser.add_argument("--save-frequency", type=int, default=4000, help="do saving checkpoint every save_frequencey step")
    parser.add_argument("--model-save-dir", type=str, default="artifacts/", help="checkpoint saving dir")
    parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path")
    main(parser.parse_args())