#!/usr/bin/env python """Training script for the ELF.""" import argparse import copy import yaml import logging import os import sys import time from functools import partial # Initialize JAX distributed BEFORE importing other JAX modules import jax try: jax.distributed.initialize() except (RuntimeError, ValueError): pass # Single-host run, or already initialized. # Ensure repo root on sys.path so imports work when run as a script REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if REPO_ROOT not in sys.path: sys.path.insert(0, REPO_ROOT) import jax.numpy as jnp import numpy as np from flax import jax_utils from flax.training.common_utils import get_metrics, shard from tqdm import tqdm from transformers import AutoTokenizer import wandb from modules.t5_encoder import get_encoder from utils.logging_utils import log_for_0 from utils.checkpoint_utils import ( save_checkpoint, load_encoder_checkpoint, load_checkpoint, find_latest_checkpoint, ) from utils.train_utils import ( TrainState, prefetch_to_device, get_optimizer, create_learning_rate_fn, ) from generation import run_generation from configs.config import load_config_from_yaml, apply_config_overrides, load_sampling_configs, SamplingConfig from modules.model import ELF_models from utils.data_utils import get_dataloader, prepare_batch, load_dataset, get_pad_token_id from train_step import train_step # Logging: no timestamps; suppress noisy checkpoint loggers; unbuffered stdout logging.basicConfig( format="%(levelname)s - %(name)s - %(message)s", handlers=[logging.StreamHandler(sys.stdout)], level=logging.INFO, force=True, ) logger = logging.getLogger(__name__) for _name in ("absl", "orbax", "tensorstore", "flax.training.checkpoints"): logging.getLogger(_name).setLevel(logging.ERROR) sys.stdout.reconfigure(line_buffering=True) def parse_args(): parser = argparse.ArgumentParser(description="Train ELF Diffusion Model (JAX).") parser.add_argument("--config", type=str, default=None, help="Path to a YAML config file to override defaults.") parser.add_argument( "--config_override", action="append", default=[], help="Override config values (field_name=value). Can be specified multiple times.", ) return parser.parse_args() # ============================================ # Main Training Loop # ============================================ def run_training(config): # Print configuration log_for_0("=" * 60) log_for_0("ELF Diffusion Model Training (JAX/Flax)") log_for_0("=" * 60) log_for_0(f"Model: {config.model}") log_for_0(f"Encoder Model: {config.encoder_model_name}") log_for_0(f"Encoder Checkpoint: {config.encoder_checkpoint}") log_for_0(f"Data: {config.data_path}") log_for_0(f"Max sequence length: {config.max_length}") log_for_0(f"Output dir: {config.output_dir}") log_for_0(f"HF Repo ID: {config.hf_repo_id}") log_for_0(f"Batch size per device: {config.batch_size}") log_for_0(f"Number of epochs: {config.epochs}") log_for_0(f"JAX devices: {jax.device_count()}") log_for_0(f"JAX backend: {jax.default_backend()}") log_for_0("=" * 60) # Initialize wandb if config.use_wandb and jax.process_index() == 0: wandb_config = {k: getattr(config, k) for k in dir(config) if not k.startswith("_")} wandb_tags = config.wandb_tag.split(",") if config.wandb_tag else None wandb.init( project=config.wandb_project, entity=config.wandb_entity, name=config.wandb_run_name, id=config.wandb_run_name, resume=config.wandb_resume, tags=wandb_tags, config=wandb_config, dir="/tmp", settings=wandb.Settings(start_method="thread"), ) resume_suffix = f" (resume={config.wandb_resume}, id={config.wandb_run_name})" log_for_0(f"Wandb initialized: {wandb.run.url}{resume_suffix}") rng = jax.random.PRNGKey(config.seed) log_for_0("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name or config.encoder_model_name) pad_token_id = get_pad_token_id(tokenizer, config.pad_token) log_for_0(f"Using {'EOS' if config.pad_token == 'eos' else 'PAD'} token for padding: {pad_token_id}") train_dataset, eval_dataset = load_dataset(config) # ============================================ # Load frozen encoder # ============================================ # Always instantiate the model to get d_model (architecture) and model.apply (static arg). log_for_0(f"Loading Encoder config: {config.encoder_model_name}...") encoder_config, encoder_model, _ = get_encoder(config.encoder_model_name, jnp.float32) encoder_params = load_encoder_checkpoint(config.encoder_checkpoint) log_for_0("Encoder weights loaded.") encoder_params = jax_utils.replicate(encoder_params) log_for_0(f"Encoder d_model: {encoder_config.d_model}") # ============================================ # Create ELF Model # ============================================ log_for_0(f"Creating {config.model} model...") model_fn = ELF_models[config.model] rng, init_rng, dropout_rng = jax.random.split(rng, 3) # Dummy inputs for initialization (2x dim if self_cond_prob > 0 to init self_cond_proj layer) max_length = config.max_length input_dim = 2 * encoder_config.d_model if config.self_cond_prob > 0 else encoder_config.d_model dummy_x = jnp.ones((1, max_length, input_dim)) dummy_t = jnp.ones((1,)) dummy_self_cond_cfg_scale = jnp.ones((1,)) if config.num_self_cond_cfg_tokens > 0 else None log_for_0(f"Dummy shapes: x={dummy_x.shape}, t={dummy_t.shape}") # Use the full tokenizer length for CE heads; tokenizer.vocab_size can exclude # added special tokens that still appear in tokenized Qwen targets. try: vocab_size = len(tokenizer) except TypeError: vocab_size = tokenizer.vocab_size log_for_0(f"Tokenizer vocab: CE head={vocab_size}") model = model_fn( text_encoder_dim=encoder_config.d_model, max_length=max_length, attn_drop=config.attn_dropout, proj_drop=config.proj_dropout, num_time_tokens=config.num_time_tokens, num_self_cond_cfg_tokens=config.num_self_cond_cfg_tokens, vocab_size=vocab_size, num_model_mode_tokens=config.num_model_mode_tokens, bottleneck_dim=config.bottleneck_dim, ) log_for_0("Initializing ELF model...") init_args = dict( x=dummy_x, t=dummy_t, deterministic=True, self_cond_cfg_scale=dummy_self_cond_cfg_scale, ) elf_params = model.init(init_rng, **init_args) log_for_0("\n" + model.tabulate(init_rng, **init_args)) total_params = sum(x.size for x in jax.tree_util.tree_leaves(elf_params)) log_for_0(f"ELF parameters: {total_params:,}") num_devices = jax.device_count() num_local_devices = jax.local_device_count() num_hosts = jax.process_count() if config.global_batch_size is not None: log_for_0(f"Using global batch size: {config.global_batch_size}") total_batch_size = config.global_batch_size local_batch_size = total_batch_size // num_hosts config.batch_size = local_batch_size elif config.batch_size is not None: log_for_0(f"Using batch size per device: {config.batch_size}") total_batch_size = config.batch_size * num_devices local_batch_size = config.batch_size * num_local_devices config.global_batch_size = total_batch_size else: raise ValueError("Either global_batch_size or batch_size must be specified") steps_per_epoch = len(train_dataset) // total_batch_size num_train_steps = steps_per_epoch * config.epochs if config.warmup_steps >= 0: num_warmup_steps = config.warmup_steps elif config.warmup_epochs is not None: num_warmup_steps = int(config.warmup_epochs * steps_per_epoch) else: num_warmup_steps = 0 # Gradient accumulation: LR schedule is parameterized in optimizer steps grad_accum_steps = config.grad_accum_steps num_optimizer_steps = num_train_steps // grad_accum_steps num_warmup_optimizer_steps = num_warmup_steps // grad_accum_steps # Effective learning rate (scaled with effective batch size, including grad accum) if config.lr is None or config.lr <= 0: if config.lr is not None: log_for_0(f"Configured lr={config.lr} is non-positive; recomputing from blr={config.blr}") config.lr = config.blr * (total_batch_size * grad_accum_steps) / 256 log_for_0( f"Hosts={num_hosts}, local_devices={num_local_devices}, total_devices={num_devices} | " f"batch local={local_batch_size}, total={total_batch_size} | " f"steps/epoch={steps_per_epoch}, total_train={num_train_steps}, warmup={num_warmup_steps}, lr={config.lr:.2e}" ) if grad_accum_steps > 1: log_for_0( f"Grad accum={grad_accum_steps}, effective batch={total_batch_size * grad_accum_steps}, " f"optimizer steps={num_optimizer_steps}" ) lr_schedule = create_learning_rate_fn( num_train_steps=num_optimizer_steps, num_warmup_steps=num_warmup_optimizer_steps, learning_rate=config.lr, schedule=config.lr_schedule, min_lr=config.min_lr, ) optimizer = get_optimizer(config, lr_schedule, grad_accum_steps=grad_accum_steps) state = TrainState.create( apply_fn=model.apply, params=elf_params["params"], tx=optimizer, dropout_rng=dropout_rng, ema_params1=copy.deepcopy(elf_params["params"]), ) total_trainable = sum(x.size for x in jax.tree_util.tree_leaves(state.params)) log_for_0(f"Total trainable parameters: {total_trainable:,}") # Auto-resume: if no explicit resume path, check output_dir for existing checkpoints if not config.resume: auto_ckpt = find_latest_checkpoint(config.output_dir) if auto_ckpt: config.resume = config.output_dir log_for_0(f"Auto-resuming from {auto_ckpt}") start_epoch, resume_step = 0, 0 resume_epoch_fractional = 0.0 # Fractional epoch for save-point tracking if config.resume: try: ckpt_path = config.resume if "checkpoint_" not in ckpt_path: ckpt_path = find_latest_checkpoint(ckpt_path) or ckpt_path state, resume_step = load_checkpoint(ckpt_path, state) resume_epoch_fractional = float(state.epoch) start_epoch = int(state.epoch) log_for_0(f"Resumed from step {resume_step} (epoch {resume_epoch_fractional:.2f})") except Exception as e: log_for_0(f"Error loading checkpoint: {e}") log_for_0("Continuing training from scratch") state = jax_utils.replicate(state) p_train_step = jax.pmap( partial(train_step, encoder_apply_fn=encoder_model.apply, config=config), axis_name="batch", donate_argnums=(0,), ) os.makedirs(config.output_dir, exist_ok=True) config_dict = { k: ([vars(sc) for sc in v] if isinstance(v, list) and v and isinstance(v[0], SamplingConfig) else v) for k, v in vars(config).items() } config_path = os.path.join(config.output_dir, 'config.yml') with open(config_path, 'w') as f: yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False) log_for_0(f"Config saved to {config_path}") train_dataloader = get_dataloader( train_dataset, batch_size=local_batch_size, shuffle=True, num_workers=config.num_workers, drop_last=True, max_seq_length=config.max_length, pad_token_id=pad_token_id, max_input_seq_length=config.max_input_length, ) # ============================================ # Checkpoint and Evaluation Schedule (Epoch-based) # ============================================ log_for_0("\n" + "=" * 60) log_for_0("Checkpoint and Evaluation Schedule") log_for_0("=" * 60) log_for_0( f"Steps/epoch={steps_per_epoch}, epochs={config.epochs}, total={steps_per_epoch * config.epochs} | " f"save every {config.save_freq} epoch(s), eval every {config.eval_freq} epoch(s)" ) if config.sampling_configs_path: config.sampling_configs = load_sampling_configs(config.sampling_configs_path) log_for_0(f"Sampling configs: {len(config.sampling_configs)} config(s)") # ============================================ # Training Loop # ============================================ log_for_0("\n" + "=" * 60) log_for_0("Starting Training") log_for_0("=" * 60) if resume_step > 0: global_step = resume_step # Skip already-processed batches within the current epoch on resume steps_to_skip_in_epoch = resume_step - start_epoch * steps_per_epoch else: global_step = start_epoch * steps_per_epoch steps_to_skip_in_epoch = 0 last_log_step = global_step train_metrics = [] last_log_time = time.time() # Track last save point for fractional save_freq; use fractional epoch from # checkpoint to avoid re-saving immediately after resume. last_save_epoch = resume_epoch_fractional if resume_step > 0 else float(start_epoch) for epoch in range(start_epoch, config.epochs): log_for_0(f"\nEpoch {epoch + 1}/{config.epochs}") # Free device buffers from previous epoch before allocating new ones, to avoid # transient OOM at epoch boundaries. if epoch > start_epoch: del train_loader, train_iterator train_metrics = [] jax.clear_caches() train_dataloader.sampler.set_epoch(epoch) train_iterator = iter(train_dataloader) train_loader = prefetch_to_device(train_iterator, size=4) initial_pbar = (resume_step - start_epoch * steps_per_epoch) if (epoch == start_epoch and resume_step > 0) else 0 epoch_pbar = tqdm( total=steps_per_epoch, desc=f"Epoch {epoch + 1}", initial=initial_pbar, mininterval=1.0, disable=jax.process_index() != 0, ) for step_in_epoch, batch in enumerate(train_loader): is_first_step = step_in_epoch == 0 and epoch == start_epoch if is_first_step: log_for_0("Performing initial training step, this may take longer...") # Skip already-processed batches when resuming mid-epoch if epoch == start_epoch and step_in_epoch < steps_to_skip_in_epoch: continue rng, batch_rng = jax.random.split(rng, 2) batch = prepare_batch(batch, config, rng=batch_rng) batch = {k: v for k, v in batch.items() if isinstance(v, (np.ndarray, jnp.ndarray))} batch = shard(batch) state, metrics = p_train_step(state, encoder_params, batch=batch) # Sync only on first step to measure XLA compilation time; # float() on the loss below already forces a device-to-host sync. if is_first_step: jax.tree_util.tree_map(lambda x: x.block_until_ready(), metrics) log_for_0("First training step (XLA compilation + execution) completed...") train_metrics.append(metrics) global_step += 1 epoch_pbar.update(1) if global_step % config.log_freq == 0: jax.tree_util.tree_map(lambda x: x.block_until_ready(), state.params) gathered = get_metrics(train_metrics) avg_loss = float(jnp.mean(gathered["loss"])) avg_l2_loss = float(jnp.mean(gathered["l2_loss"])) avg_ce_loss = float(jnp.mean(gathered["ce_loss"])) now = time.time() steps_per_sec = (global_step - last_log_step) / max(now - last_log_time, 1e-8) current_lr = lr_schedule((global_step - 1) // grad_accum_steps) postfix_dict = { "step": f"{global_step}", "loss": f"{avg_loss:.4f}", "l2": f"{avg_l2_loss:.4f}", "ce": f"{avg_ce_loss:.4f}", "sps": f"{steps_per_sec:.1f}", "lr": f"{current_lr:.2e}", } log_for_0(postfix_dict) epoch_pbar.set_postfix(**postfix_dict) if jax.process_index() == 0: tqdm.write( f"INFO - engine - Step {global_step}: loss={avg_loss:.4f}, " f"l2={avg_l2_loss:.4f}, ce={avg_ce_loss:.4f}, " f"lr={current_lr:.2e}, steps/sec={steps_per_sec:.2f}" ) if config.use_wandb: current_epoch_progress = epoch + (step_in_epoch + 1) / steps_per_epoch try: wandb.log({ "train_loss": avg_loss, "train_l2_loss": avg_l2_loss, "train_ce_loss": avg_ce_loss, "lr": current_lr, "epoch": current_epoch_progress, "step": global_step, }, step=global_step) except Exception: pass train_metrics = [] last_log_step = global_step last_log_time = now # Intra-epoch checkpoint saving (fractional save_freq, e.g., 0.1 epoch) if 0 < config.save_freq < 1: progress = epoch + (global_step - epoch * steps_per_epoch) / steps_per_epoch if progress - last_save_epoch >= config.save_freq: save_checkpoint(state, config.output_dir, global_step, hf_repo_id=config.hf_repo_id) log_for_0(f"Saved checkpoint at epoch {progress:.2f} (step {global_step})") last_save_epoch = progress epoch_pbar.close() current_epoch = epoch + 1 state = jax_utils.replicate(jax_utils.unreplicate(state).replace(epoch=current_epoch)) if config.save_freq >= 1 and current_epoch % config.save_freq == 0: save_checkpoint(state, config.output_dir, global_step, hf_repo_id=config.hf_repo_id) log_for_0(f"Saved checkpoint at epoch {current_epoch} (step {global_step})") if config.eval_freq >= 1 and current_epoch % config.eval_freq == 0: rng = run_generation( state=state, encoder_params=encoder_params, encoder_apply_fn=encoder_model.apply, eval_dataset=eval_dataset, tokenizer=tokenizer, config=config, rng=rng, local_batch_size=local_batch_size, ) last_log_step = global_step last_log_time = time.time() log_for_0("\n" + "=" * 60) log_for_0("Final Generation") log_for_0("=" * 60) save_checkpoint(state, config.output_dir, global_step, hf_repo_id=config.hf_repo_id) log_for_0(f"Final checkpoint saved to {config.output_dir}") if config.use_wandb and jax.process_index() == 0: wandb.finish() def main(): """CLI entry point: parse args, load config, then run training.""" args = parse_args() config = load_config_from_yaml(args.config) if args.config_override: config = apply_config_overrides(config, args.config_override) log_for_0(f"Applied {len(args.config_override)} config override(s)") run_training(config) if __name__ == "__main__": main()