# coding=utf-8 # Copyright 2024 The Google Research Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """This file contains function that coordinates the training and evaluation of the model.""" import functools import math import os from absl import logging from clu import metric_writers from flax import jax_utils from flax import linen as nn from flax.training import checkpoints import jax from jax import random import jax.numpy as jnp import numpy as np import tensorflow as tf import wandb from train import data from train import evaluater from train import model from train import trainer def log_hyperparams_tb( config, model_config, initial_variables, tf_summary_writer ): """Log hyperparameters to TensorBoard. Args: config: experiment's ConfigDict model_config: model's ConfigDict initial_variables: initial hyperparameter values tf_summary_writer: SummaryWriter object. Returns: The SummaryWriter object and the config. """ # Calculate the total number of model parameters config.num_model_parameters = sum( x.size for x in jax.tree_util.tree_leaves(initial_variables) ) # Convert hyperparameters to tensors config_hyperparameters = [ tf.convert_to_tensor([k, str(v)]) for k, v in config.items() ] model_config_hyperparameters = [ tf.convert_to_tensor([k, str(v)]) for k, v in model_config.__dict__.items() ] # Log model hyperparameters to TensorBoard with tf_summary_writer.as_default(): tf.summary.text( "Model hyperparameters", tf.stack(model_config_hyperparameters), step=0 ) tf.summary.text( "Config hyperparameters", tf.stack(config_hyperparameters), step=0 ) return tf_summary_writer, config def train_and_evaluate(config, workdir): """The training and evaluation loops for the model. Args: config: experiment's config dictionary. workdir: directory to use for logging. """ # Orbax checkpointing requires an absolute path. workdir = os.path.abspath(workdir) logging.info("Creating training and evaluator dataset iterator") curriculum = data.CurriculumState( stage=int(getattr(config, "curriculum_start_stage", 1)), max_stage=int(getattr(config, "curriculum_max_stage", 6))) train_data_iter = data.create_iter( config, config.minibatch_size, train=True, curriculum=curriculum) eval_data_iter = data.create_iter(config, config.minibatch_size, train=False) logging.info("Finished creating training dataset iterator") model_config = model.TransformerConfig( dtype=config.dtype, vocab_size=config.vocab_size, seq_len=config.seq_len, num_heads=config.num_heads, num_layers=config.num_layers, emb_dim=config.emb_dim, qkv_dim=config.qkv_dim, mlp_dim=config.mlp_dim, dropout_rate=config.dropout_rate, attention_dropout_rate=config.attention_dropout_rate, deterministic=False, num_latent_slots=int(config.num_latent_slots), inject_latents=bool(int(getattr(config, "recurrent_latent", 1))), ) logging.info("train_config: %s", str(model_config.__dict__)) print(str(model_config.__dict__), flush=True) rng = jax.random.PRNGKey(config.seed) rng, init_rng, inference_rng = random.split(rng, num=3) # Initialize the model and get initial variables. Dummy latent arguments # are provided so the latent projector parameters are created at init. rng, dropout_rng = jax.random.split(rng) input_shape = (config.minibatch_size, config.seq_len) net = model.TransformerLMHeadModel(model_config) rng_keys = {"params": init_rng, "dropout": dropout_rng} K = int(config.num_latent_slots) dummy_latents = jnp.zeros( (config.minibatch_size, K, config.emb_dim), model_config.dtype) dummy_positions = jnp.zeros((config.minibatch_size, K), jnp.int32) dummy_active = jnp.zeros((config.minibatch_size, K), bool) sample_out, initial_variables = jax.jit( net.init_with_output )(rng_keys, jnp.ones(input_shape, jnp.int32), dummy_latents, dummy_positions, dummy_active) state, lr_scheduler_fn = trainer.get_state(config, net, initial_variables) # Resume-and-extend support: when resuming, start the training loop at the # restored optimizer step (not 0) so a larger config.max_steps continues the # cosine LR schedule cleanly instead of re-running steps or overshooting. start_step = 0 if config.resume_training: state = checkpoints.restore_checkpoint(config.ckpt_loc, state) start_step = int(state.step) print("----------Restored model from", config.ckpt_loc, f"at step {start_step}-----------") writer = metric_writers.create_default_writer( workdir, asynchronous=False, just_logging=(jax.process_index() > 0)) tf_summary_writer = tf.summary.create_file_writer(workdir) logging.info("config: %s", str(config.__dict__)) state = jax_utils.replicate(state) dropout_rngs = jax.random.split(rng, jax.local_device_count()) def make_p_train_step(num_passes): return jax.pmap( functools.partial( trainer.train_step, config=model_config, hyperparams=config, learning_rate_fn=lr_scheduler_fn, num_passes=num_passes), axis_name="batch", donate_argnums=(0,)) # num_passes = latent recurrence depth for this stage, capped by the number # of latent slots (K=0 -> 0 passes -> no-latent control baseline). # When recurrence is disabled (control), force 0 passes so no hidden state # is ever fed back: the slots become static, independent per-stage readouts. recurrent = bool(int(getattr(config, "recurrent_latent", 1))) passes_per_stage = int(getattr(config, "passes_per_stage", 1)) def passes_for(stage): return min(passes_per_stage * stage, K) if recurrent else 0 p_train_step = make_p_train_step(passes_for(curriculum.stage)) p_eval_step = jax.pmap(functools.partial(evaluater.eval_step, config=model_config.replace(deterministic=True)), axis_name="batch") hooks, report_progress, train_metrics = trainer.get_metrics_report_progress( config, workdir, writer) tf_summary_writer, config = log_hyperparams_tb( config, model_config, initial_variables, tf_summary_writer ) promote_threshold = float(getattr(config, "promote_acc_threshold", 0.85)) promote_loc_threshold = float(getattr(config, "promote_loc_threshold", 0.70)) promote_patience = int(getattr(config, "promote_patience_steps", 8000)) min_stage_steps = int(getattr(config, "min_stage_steps", 2000)) plateau_steps = int(getattr(config, "plateau_steps", 0)) plateau_delta = float(getattr(config, "plateau_delta", 0.005)) instance_mode = bool(getattr(config, "instance_dir", None)) per_stage_inset = {} # Best frontier-depth accuracy seen in the current stage, and the step it # was last improved: the plateau detector's state. stage_best_acc = -1.0 stage_best_step = start_step ckpt_keep = int(getattr(config, "ckpt_keep", 100)) # Protected directory for per-stage checkpoints (never rolled off). stage_ckpt_dir = os.path.join(workdir, "stage_ckpts") stage_started_at = start_step with metric_writers.ensure_flushes(writer): for step in range(start_step, config.max_steps): if step%10000 == 0: print("Step:", step, flush=True) state, metrics = trainer.train_one_step(p_train_step, config, state, step, dropout_rngs, train_data_iter) for h in hooks: h(step) if math.isnan(metrics["loss"][0]): print("The loss function became nan: This might be due to the choice of hyperparameters.") break if step % config.eval_every_steps == 0: eval_metrics = evaluater.get_eval_metrics( state, eval_data_iter, p_eval_step, config) per_level = eval_metrics.pop("per_level_acc") per_slot = eval_metrics.pop("per_slot_acc", {}) per_slot_ch = eval_metrics.pop("per_slot_acc_changed", {}) per_stage_inset = eval_metrics.pop("per_stage_inset_acc", {}) per_bin = eval_metrics.pop("per_bin_acc", {}) def _m(key): v = eval_metrics.get(key, []) return round(float(np.mean(v)), 4) if len(v) else -1.0 print(step, "stage", curriculum.stage, "loss", round(float(metrics["loss"].mean()), 4), "ce", round(float(metrics["ce_loss"].mean()), 4), "aux_bce", round(float(metrics["aux_loss"].mean()), 4), "| val_acc", _m("acc"), "loc_acc", _m("loc_acc"), "val|loc", _m("val_given_loc_acc"), "| cand_bit_acc", _m("cand_bit_acc"), "cand_set_acc", _m("cand_set_acc"), "cand_set_chg", _m("cand_set_acc_changed"), "| inset", round( float(per_stage_inset.get(curriculum.stage, -1.0)), 4), "promote_need inset>={:.2f} loc>={:.2f}".format( promote_threshold, promote_loc_threshold), "| depth_acc", round( float(per_slot_ch.get(curriculum.stage, -1.0)), 4), flush=True) with tf_summary_writer.as_default(): tf.summary.scalar("loss", metrics["loss"].mean(), step=step) tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step) tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step) tf.summary.scalar( "learning rate", metrics["learning_rate"].mean(), step=step ) tf.summary.scalar("curriculum_stage", curriculum.stage, step=step) log_dict = {'loss': metrics["loss"].mean(), 'learning rate': metrics["learning_rate"].mean()} for key in eval_metrics.keys(): vals = eval_metrics[key] if not vals: continue tf.summary.scalar( "eval_" + key, np.array(vals).mean(), step=step ) log_dict[ "eval_" + key ] = np.array(vals).mean() for lvl, v in per_level.items(): if v >= 0: tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step) log_dict[f"eval_acc_level_{lvl}"] = v for s, v in per_slot.items(): if v >= 0: tf.summary.scalar(f"eval_cand_depth_{s}", v, step=step) log_dict[f"eval_cand_depth_{s}"] = v for s, v in per_slot_ch.items(): if v >= 0: tf.summary.scalar(f"eval_cand_depth_chg_{s}", v, step=step) if config.use_wandb: wandb.log(log_dict, step=step) # ---- Curriculum promotion (ATC-style: threshold + patience) ---- # Two curricula, one frontier rule. With latent slots the # frontier is the DEEPEST active slot, i.e. the newest wave # snapshot this stage introduced, scored on changed cells only # (unchanged cells are copies of slot j-1 and stay correct for a # head that learned nothing new). With the round-count data # curriculum the frontier is the newest round-BIN, i.e. the # longest-chain puzzles this stage admitted. rounds_curric = str( getattr(config, "data_curriculum", "none")) == "rounds" has_frontier = int(config.num_latent_slots) > 0 or rounds_curric if has_frontier and curriculum.stage < curriculum.max_stage: t = curriculum.stage loc_now = _m("loc_acc") if instance_mode: # Inset = "emitted digit is in the stage-t set", scored # only on location-correct cells. Require BOTH: the # solver-order locations, and values inside the set. # Otherwise a high inset on ~2% loc_acc promotes from # noise. Plateau is also blocked until location is up, # so we cannot skip stage 1 the way the old patience # promotions did at 0.32-0.46. inset_now = per_stage_inset.get(t, -1.0) frontier_acc = float(inset_now) if inset_now >= 0 else -1.0 elif int(config.num_latent_slots) > 0: frontier_acc = per_slot_ch.get(t, -1.0) if frontier_acc < 0: frontier_acc = per_slot.get(t, -1.0) else: frontier_acc = per_bin.get(t, -1.0) if frontier_acc < 0: # The frontier bin holds as little as 6.5% of the # corpus, so a single eval can miss it. Fall back to # the pooled accuracy over everything unlocked. seen = [v for b, v in per_bin.items() if b <= t and v >= 0] frontier_acc = (float(np.mean(seen)) if seen else -1.0) steps_in_stage = step - stage_started_at # A negative accuracy means "not measured this eval". Never # let that count as evidence: it must not reset the plateau # tracker, and it must not satisfy the threshold or plateau # rule. Only the hard patience cap can fire without a # measurement. measured = frontier_acc >= 0 if measured and frontier_acc > stage_best_acc + plateau_delta: stage_best_acc = frontier_acc stage_best_step = step loc_ready = (not instance_mode) or loc_now >= promote_loc_threshold hit_threshold = (measured and loc_ready and frontier_acc >= promote_threshold) stalled = (measured and loc_ready and plateau_steps > 0 and (step - stage_best_step) >= plateau_steps) patience_over = steps_in_stage >= promote_patience if steps_in_stage >= min_stage_steps and ( hit_threshold or stalled or patience_over): reason = ("threshold" if hit_threshold else "plateau" if stalled else "patience") curriculum.stage += 1 stage_started_at = step stage_best_acc = -1.0 stage_best_step = step p_train_step = make_p_train_step(passes_for(curriculum.stage)) what = ("round-bin" if int(config.num_latent_slots) == 0 else "depth") print(f"[curriculum] step {step}: promote to stage " f"{curriculum.stage} ({reason}; graduated {what} " f"{t} acc={frontier_acc:.3f} after " f"{steps_in_stage} steps); " f"latent passes={curriculum.stage}, " f"pool/snapshots now 1..{curriculum.stage}", flush=True) if config.save_checkpoint: unrep_state = jax_utils.unreplicate(state) # Rolling checkpoint in the main workdir. checkpoints.save_checkpoint_multiprocess( workdir, unrep_state, step, keep=ckpt_keep, overwrite=True) # Protected copy: the model *entering* this stage, # kept permanently under stage_ckpts/ (never rolled # off), so every stage's checkpoint survives. checkpoints.save_checkpoint_multiprocess( stage_ckpt_dir, unrep_state, step, keep=100, overwrite=True, prefix=f"stage{curriculum.stage}_") if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0: checkpoints.save_checkpoint_multiprocess( workdir, jax_utils.unreplicate(state), step, keep=ckpt_keep, overwrite=True ) # Final checkpoint at the end of training if config.save_checkpoint: checkpoints.save_checkpoint_multiprocess( workdir, jax_utils.unreplicate(state), config.max_steps, keep=ckpt_keep, overwrite=True)