src/imagenet/train_model.py

import argparse
import os
import jax
import flax
import copy
import json
import optax
import torch
import wandb
import numpy as np
from tqdm import tqdm
import jax.numpy as jnp
from flax import linen as nn
from flax.jax_utils import replicate, unreplicate
from flax.training import checkpoints, train_state
from flax.core.frozen_dict import freeze, unfreeze
from flax.traverse_util import flatten_dict, unflatten_dict
from transformers.models.vit.modeling_flax_vit import ViTConfig, FlaxViTForImageClassification
from flax.training.common_utils import get_metrics, onehot, shard
from lmc_model import  LMCFlaxViTForImageClassification, print_model, print_model_with_prefix
from datasets import build_dataset
import multiprocessing as mp
from jax import debug
from pprint import pprint
from typing import Any, Dict, List
import shutil
mp.set_start_method("spawn", force=True)
os.environ["WANDB_API_KEY"] = "a31cc923ed06cb0ea30db501a39433d07f7789a9"
def aggregate_metrics(metrics_list):
    return {
        key: jnp.mean(jnp.array([m[key] for m in metrics_list])).item()
        for key in metrics_list[0]
    }
# ---------- Dataset Loader ----------
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 imagenet_data_loader(args):
    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
    dataset_val, _ = build_dataset(is_train=False, args=args)
    sampler_train = torch.utils.data.RandomSampler(dataset_train)
    sampler_val = torch.utils.data.SequentialSampler(dataset_val)
    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train, batch_size=args.batch_size,
        num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True
    )
    data_loader_val = torch.utils.data.DataLoader(
        dataset_val, sampler=sampler_val, batch_size=args.batch_size,
        num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False
    )
    return data_loader_train, data_loader_val
# fmt: on
def prepare_image_batch(images:torch.Tensor,labels:torch.Tensor) -> Dict[str, Any]:
    images, labels =  jnp.array(images),jnp.array(labels)
    return {'images': shard(images),'labels': shard(labels)}
def accuracy(logits, labels, topk=(1,)):
    maxk = max(topk)
    batch_size = labels.shape[0]
    topk_preds = jnp.argsort(logits, axis=-1)[:, -maxk:][:, ::-1]  # Top-k predictions
    res = []
    for k in topk:
        correct = (topk_preds[:, :k] == labels[:, None])
        correct = jnp.any(correct, axis=1)
        correct = jnp.sum(correct)
        res.append(100.0 * correct / batch_size)
    return res  # list of [acc@1, acc@5]
def main(args: argparse.Namespace):
    # --- Seeds & RNG ---
    wandb.init(
        project=args.wandb_project,
        entity=args.wandb_entity,
        group=args.wandb_group,
        id=args.wandb_id,
        name= f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-shared{args.num_shared_experts}-seed{args.seed}",
        save_code=True
    )
    wandb.config = dict(vars(args))
    save_path = os.path.join(args.save_dir,f"lr{args.lr}-{args.position_embeddings}-epochs{args.epochs}-batch{args.batch_size}-seed{args.seed}")
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    rng = jax.random.PRNGKey(args.seed)
    # --- Prepare Data Loader ---
    train_loader, val_loader = imagenet_data_loader(args)
    # --- Load pretrained model ---
    config = ViTConfig()
    config.hidden_size = args.hidden_size
    config.num_hidden_layers = args.num_hidden_layers
    config.num_attention_heads = args.num_attention_heads
    config.intermediate_size = args.intermediate_size
    config.patch_size = args.patch_size
    config.image_size = args.input_size
    config.num_labels = args.num_labels
    config.position_embeddings = args.position_embeddings
    config.rotary_value = args.rotary_value
    config.num_shared_experts = args.num_shared_experts
    config.num_routed_experts = args.num_routed_experts
    config.topk = args.topk
    config.routed_scaling_factor = args.routed_scaling_factor
    config.lmc_layer_indices = args.lmc_layer_indices
    model = LMCFlaxViTForImageClassification(
        config,input_shape=(1,config.image_size, config.image_size, config.num_channels),seed=args.seed,dtype=jnp.dtype(args.dtype),
    )
    model.config.save_pretrained(save_path)
    num_global_steps = len(train_loader)*args.epochs 
    num_warmup_steps = len(train_loader)*args.warmup_epochs
    lr_schedule =optax.warmup_cosine_decay_schedule(
        init_value=args.warmup_lr,
        peak_value=args.lr,
        warmup_steps=num_warmup_steps,
        decay_steps=num_global_steps,
        end_value=args.min_lr,
    )
    tx = optax.adamw(
        learning_rate=lr_schedule,
        b1=args.adamw_beta1,
        b2=args.adamw_beta2,
        eps=args.adamw_eps,
        weight_decay=args.weight_decay,
    )
    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=tx)
    if args.restore_checkpoint_path:
        state = checkpoints.restore_checkpoint(args.restore_checkpoint_path, state)
        print(f"train state restored from {args.restore_checkpoint_path}")
        print(f"skip trian step to {state.step}")
    latest_global_step = state.step
    curr_epoch = latest_global_step//len(train_loader)
    state = replicate(state)
    def train_step(state, batch, rng):
        dropout_rng, new_dropout_rng = jax.random.split(rng)
        def loss_fn(params):
            outputs = state.apply_fn(params=params,pixel_values=batch["images"],train=True,dropout_rng=dropout_rng,)
            logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
            loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean()
            return loss, logits
        (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
        grads = jax.lax.pmean(grads, axis_name="batch")
        state = state.apply_gradients(grads=grads)
        acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5))
        metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,"learning_rate": lr_schedule(state.step),}
        metrics = jax.lax.pmean(metrics, axis_name="batch")
        return state, metrics, new_dropout_rng
    def eval_step(state, batch):
        outputs = state.apply_fn(params=state.params,pixel_values=batch["images"],train=False,)
        logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
        loss = optax.softmax_cross_entropy_with_integer_labels(logits, batch["labels"]).mean()
        acc1, acc5 = accuracy(logits, batch["labels"], topk=(1, 5))
        metrics = {"loss": loss,"acc1": acc1,"acc5": acc5,}
        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")
    rng = jax.random.PRNGKey(args.seed)
    global_step = latest_global_step
    best_val_acc1 = 0
    print("Starting training...")
    print(f"JAX devices: {jax.devices()}")
    print(f"Using {jax.local_device_count()} devices")
    for epoch in range(curr_epoch,args.epochs):
        print(f"Epoch {epoch}")
        dropout_rngs = jax.random.split(rng, jax.local_device_count())
        train_metrics_stack = []
        pbar = tqdm(enumerate(train_loader), desc="Training", leave=False)
        for batch_idx, (images, labels) in pbar:
            # Prepare and shard batch
            batch = prepare_image_batch(images,labels)
            # Run train step
            state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
            train_metrics = unreplicate(train_metrics)
            train_metrics = jax.tree_util.tree_map(jnp.mean, train_metrics)
            loss, acc1, acc5 = (float(train_metrics["loss"]),float(train_metrics["acc1"]),float(train_metrics["acc5"]),)            
            curr_lr = float(lr_schedule(global_step))
            if global_step % args.wandb_logging_frequency == 0:
                wandb.log({"train/loss": loss,"train/acc1": acc1,"train/acc5": acc5,"lr": curr_lr,}, step=global_step)
            pbar.set_postfix({"train/loss": f"{loss:.3f}","train/acc1": f"{acc1:.3f}", "train/acc5": f"{acc5:.3f}","lr": f"{curr_lr:.2e}",})
            global_step += 1
        eval_results = []
        pbar = tqdm(enumerate(val_loader), desc="Evaluating", leave=False)
        for batch_idx, (images, labels) in pbar:
            batch = prepare_image_batch(images, labels)
            eval_metric = parallel_eval_step(state, batch)
            eval_results.append(eval_metric)            
        # Compute mean metrics across all eval batches
        eval_metrics = get_metrics(eval_results)
        eval_metrics = unreplicate(eval_metrics)
        eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
        val_loss, val_acc1, val_acc5 = float(eval_metrics["loss"]), float(eval_metrics["acc1"]), float(eval_metrics["acc5"])
        wandb.log({"val/loss": val_loss,"val/acc1": val_acc1, "val/acc5" : val_acc5}, step=global_step)
        # Save best model
        if best_val_acc1 <= val_acc1:
            best_val_acc1 = val_acc1
            # model.params = unreplicate(state).params
            # best_dir = os.path.join(save_path, f"best_{global_step}")
            # model.save_pretrained(best_dir)
            # print(f"Best model saved at step {global_step}")
            # remove_old_dirs_with_prefix(save_path, "best_", global_step)
            checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="best_",keep=1)                    
        # Save last model
        # model.params = unreplicate(state).params
        # last_dir = os.path.join(save_path, f"last_{global_step}")
        # model.save_pretrained(last_dir)
        # print(f"Last model saved at step {global_step}")
        # remove_old_dirs_with_prefix(save_path, "last_", global_step)
        checkpoints.save_checkpoint(ckpt_dir=save_path,target=unreplicate(state),step=global_step,prefix="last_",keep=1)
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fine-tune ViT with MoE on Imagenet")
    # --- Model & Training Config ---
    parser.add_argument("--hidden-size", type=int, default=768, help="Dimensionality of the encoder layers and the pooler layer.")
    parser.add_argument("--num-hidden-layers", type=int, default=12, help="Number of hidden layers in the Transformer encoder.")
    parser.add_argument("--num-attention-heads", type=int, default=12, help="Number of attention heads for each attention layer in the Transformer encoder.")
    parser.add_argument("--intermediate-size", type=int, default=3072, help="Dimensionality of the intermediate (feed-forward) layer in the Transformer encoder.")
    parser.add_argument("--position-embeddings", type=str, default='sinusoidal')
    parser.add_argument("--patch-size", type=int, default=16, help="")
    parser.add_argument("--num-labels", type=int, default=1000, help="")
    parser.add_argument("--rotary_value", action="store_true",help='Whether or not apply rotary position embeddings on value layer.')
    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('--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('--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("--epochs", type=int, default=30)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER')
    parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR', help='warmup learning rate (default: 1e-6)')
    parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',help='epochs to warmup LR, if scheduler supports')
    parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')
    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("--patience", type=int, default=10, help="Early stopping patience")
    parser.add_argument("--restore-checkpoint-path", type=str, help="if you want to restart from specific checkpoint, set this arg to checkpoint path")
    # --- Data Config ---
    parser.add_argument("--data-path", type=str, required=True)
    parser.add_argument("--save-dir", type=str, required=True)
    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR10','CIFAR100', 'IMNET', 'INAT', 'INAT19'])
    parser.add_argument("--input-size", type=int, default=224)
    parser.add_argument('--num_workers', type=int, default=10)
    parser.add_argument('--pin-mem', action='store_true')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--color-jitter', type=float, default=0.4)
    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1')
    parser.add_argument('--train-interpolation', type=str, default='bicubic')
    parser.add_argument('--reprob', type=float, default=0.25)
    parser.add_argument('--remode', type=str, default='pixel')
    parser.add_argument('--recount', type=int, default=1)
    parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16", help="model datatype")
    parser.add_argument("--wandb-entity", default=None, help="wandb entity for logging")
    parser.add_argument("--wandb-group", default=None, help="wandb group for logging")
    parser.add_argument("--wandb-project", default=None, help="wandb project name for logging")
    parser.add_argument("--wandb-id", default=None, help="wandb project name for logging")
    parser.add_argument("--wandb-logging-frequency", type=int, default=100, help="do logging every logging_frequency step")
    args = parser.parse_args()
    main(args)