# 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. """Transformer LM trainer.""" import functools from clu import periodic_actions from flax.training import common_utils from flax.training import train_state import jax from jax import numpy as jnp import numpy as np import optax import ml_collections from train import model def get_state(config, net, initial_variables): """Get the train state given an experiment config, a model and initial variables. Args: config: A ConfigDict containing the configuration for the experiment. net: The model to use for training. initial_variables: The initial variables for the model. Returns: A tuple containing the train state and the learning rate schedule. """ # Learning rate schedule lr_scheduler_fn = functools.partial( lr_scheduler, learning_rate=config.learning_rate, warmup_tokens=config.warmup_tokens, final_tokens=config.max_steps, config=config, ) # Optimizer optim_fn = optax.adamw( lr_scheduler_fn, weight_decay=config.weight_decay, b1=0.9, b2=0.95 ) # Clip the gradients to prevent exploding gradients. optimizer = optax.chain(optax.clip_by_global_norm(1), optim_fn) # Initialize the train state state = train_state.TrainState.create( apply_fn=net.apply, params=initial_variables["params"], tx=optimizer ) return state, lr_scheduler_fn def lr_scheduler( n_tokens: int, learning_rate: float, warmup_tokens: int, final_tokens: int, config: ml_collections.ConfigDict, ) -> float: """Learning rate scheduler, adapted from Mikhail Grankin. The learning rate schedule is cosine decay with a warmup period. The learning rate starts at 0 and linearly increases to the given learning rate over the warmup period. After the warmup period, the learning rate decays according to a cosine schedule, with the given learning rate as the maximum value. Args: n_tokens: The number of tokens processed so far. learning_rate: The initial learning rate. warmup_tokens: The number of tokens to warm up over. final_tokens: The total number of tokens to process. config: A ConfigDict containing the configuration for the learning rate schedule. Returns: The learning rate at the given point in the schedule. """ # Decay the learning rate based on our progress. progress = (n_tokens - warmup_tokens) / max( 1, final_tokens - warmup_tokens, ) lr_mult = jnp.where( n_tokens < warmup_tokens, # Linear warmup. n_tokens / jnp.fmax(1, warmup_tokens), # Cosine learning rate decay. jnp.fmax(config.end_lr_factor, 0.5 * (1.0 + jnp.cos(np.pi * progress))), ) return learning_rate * lr_mult def get_metrics_report_progress(config, workdir, writer): """ Get the metrics for reporting progress during training. Args: config: The configuration for the experiment. workdir: The directory for storing the logs. writer: The writer object for recording the metrics. Returns: hooks: List of hooks for tracking progress. report_progress: Object for reporting progress. train_metrics: List of training metrics. """ hooks = [] # Initialize the report progress object report_progress = periodic_actions.ReportProgress( num_train_steps=config.max_steps, writer=writer) # Add metrics for profiling if the process index is 0 if jax.process_index() == 0: hooks += [report_progress, periodic_actions.Profile(logdir=workdir, num_profile_steps=5)] # Initialize the list of training metrics train_metrics = [] return hooks, report_progress, train_metrics def get_input_start_index(batch, config): inputs = jax.tree_util.tree_map(np.asarray, batch[0]) puzzles = jax.tree_util.tree_map(np.asarray, batch[1]) start_index = jax.tree_util.tree_map(np.asarray, batch[2]) levels = jax.tree_util.tree_map(np.asarray, batch[3]) cand_targets = jax.tree_util.tree_map(np.asarray, batch[4]) return inputs, puzzles, start_index, levels, cand_targets def train_one_step(p_train_step, config, state, step, dropout_rngs, train_data_iter): """ Single step of the training loop. Args: p_train_step: The training step function. config: The experiment configuration. state: The train state. step: The step number. dropout_rngs: The dropout random number generator. train_data_iter: The iterator for the train data. Returns: The updated train state and train metrics. """ with jax.profiler.StepTraceAnnotation("train", step_num=step): # Get the next batch from the iterator batch = next(train_data_iter) # Extract the inputs, start index and difficulty level from the batch inputs, _, start_index, levels, cand_targets = get_input_start_index(batch, config) # Shard across the devices inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs)) start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index)) levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels)) # Explicit reshape (not common_utils.shard) so a zero-width slot dim # (K=0 no-latent baseline: cand_targets is (bs, 0, 81)) shards without # the ambiguous -1 inference. Identical to shard() when K>0. cand_targets = np.asarray(cand_targets) _nd = jax.local_device_count() cand_targets = cand_targets.reshape( (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:]) # Run the training step state, metrics, _ = p_train_step( state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs ) return state, metrics def build_latent_state(inputs, start_index, levels, config, num_passes, apply_fn, rngs=None, trunc_passes=0): """Build the continuous latent thoughts (Coconut/ATC-style recurrence). Pass 0 seeds z_1 from the last-layer hidden at the last clue token; pass j reads the hidden at latent slot j to produce z_{j+1}. Slot j is active for an example only when j < k, with k = clip(level - 2, 1, K) (difficulty- matched latent budget). By default gradients flow through all passes (full BPTT). If trunc_passes=T>0, the forward chain is unchanged but thoughts older than the last T CoTs are stop-gradient'd, so the loss only backprops through those last T recurrent steps. Returns (latent_vals, slot_pos, active_full). """ num_slots = int(config.num_latent_slots) bs = inputs.shape[0] bidx = jnp.arange(bs) si3 = 3 * start_index.reshape(-1) # (bs,) slot_pos = si3[:, None] + jnp.arange(num_slots)[None, :] # (bs, K) if getattr(config, "cand_slot_mode", "level") == "depth": # Every slot the recurrence fills is active, uniformly over the batch. k = jnp.full((bs,), max(min(num_passes, num_slots), 1), dtype=jnp.int32) else: k = jnp.clip(levels.reshape(-1) - 2, 1, num_slots) # (bs,) active_full = jnp.arange(num_slots)[None, :] < k[:, None] # (bs, K) latent_vals = jnp.zeros((bs, num_slots, config.emb_dim), dtype=config.dtype) for j in range(num_passes): act_j = active_full & (jnp.arange(num_slots)[None, :] < j) _, hidden, _ = apply_fn(inputs, latent_vals, slot_pos, act_j, rngs) src = si3 - 1 + j # j=0: last clue token; j>0: latent slot j-1 z = hidden[bidx, src] # Truncated BPTT: keep the full forward chain, but detach CoTs # before the last T so earlier thoughts are treated as constants. if trunc_passes > 0 and j < num_passes - trunc_passes: z = jax.lax.stop_gradient(z) latent_vals = latent_vals.at[:, j].set(z.astype(config.dtype)) return latent_vals, slot_pos, active_full def train_step(state, batch, start_index, levels, cand_targets, config, hyperparams, learning_rate_fn, num_passes, dropout_rng=None, backtrack=False): """One step of the training loop. Args: state: Train state. batch: Input batch (bs, 3*81 + K) with latent placeholder slots. start_index: Number of clue cells per example. levels: Puzzle difficulty level (3..8) per example. config: Model config. hyperparams: Hyperparameter dictionary. learning_rate_fn: Learning rate function. num_passes: Number of latent recurrence passes (= curriculum stage). dropout_rng: RNG used for dropout. Returns: A new train state, train metrics, and computed model predictions. """ num_slots = int(config.num_latent_slots) # Extract inputs and labels from the batch inputs = batch[:, :-1] label = batch[:, 1:] # Update dropout_rng dropout_rng = jax.random.fold_in(dropout_rng, state.step) dropout_rng_dict = {"dropout": dropout_rng} def loss_fn(params): """Compute the loss function.""" net = model.TransformerLMHeadModel(config) def apply_fn(x, lv, lp, la, rngs): return net.apply({"params": params}, x, latent_values=lv, latent_positions=lp, latent_active=la, rngs=rngs) latent_vals, slot_pos, active_full = build_latent_state( inputs, start_index, levels, config, num_passes, apply_fn, rngs=dropout_rng_dict, trunc_passes=int(getattr(hyperparams, "bptt_trunc", 0))) pred_logits, _, cand_logits = apply_fn( inputs, latent_vals, slot_pos, active_full, dropout_rng_dict) label_one_hot = jax.nn.one_hot(label, num_classes=config.vocab_size) # The variables label_one_hot and pred_logits both are 3-dimensional tensors with # first axis corresponding to batch size, second correspondingn to sequence length # and third corresponding to the row/column/value at a particular cell assert label_one_hot.shape == pred_logits.shape, ("one hot label shape", label_one_hot.shape, label.shape, pred_logits.shape) # Calculate the cross-entropy loss along the last axis pred_logits_sol = pred_logits[:, :, :] label_one_hot_sol = label_one_hot[:, :, :] ce_loss = optax.softmax_cross_entropy( logits=pred_logits_sol[:, :, :], labels=label_one_hot_sol[:, :, :] ) # assert ce_loss.ndim == 2, ("ce_loss", ce_loss.shape) # Apply masking to the loss: supervise only the solution region, # which now starts K latent slots after the clue block. mask = np.repeat( np.arange(len(ce_loss[0])).reshape(1, -1), len(ce_loss), axis=0 ) mask = (mask >= 3 * start_index + num_slots) # Per-token mean (not per-example sum): keeps the LM CE on the same # O(1) scale as the per-digit BCE below, so aux_cand_weight~1 actually # balances the two instead of the candidate signal being swamped. ce_denom = jnp.maximum(mask.sum(), 1.0) avg_ce_loss = (ce_loss * mask).sum() / ce_denom # ---- Auxiliary multi-candidate BCE loss on the latent slots ---- # cand_targets: (bs, K, 81) int bitmask (bit d-1 set <=> digit d is a # candidate at that slot's stage; slot j already aligned to the puzzle's # k=level-2 budget in the data pipeline). Expand to (bs,K,81,9) multi-hot. aux_weight = float(getattr(hyperparams, "aux_cand_weight", 1.0)) if aux_weight == 0.0: # Instance arm: the superposition lives in the varying targets, not # in a set head. Drop the whole BCE graph so nothing but the LM CE # shapes the latents. zero = jnp.zeros((), dtype=pred_logits.dtype) return avg_ce_loss, (pred_logits, avg_ce_loss, zero) bits = jnp.arange(9) cand_multi_hot = ((cand_targets[..., None].astype(jnp.int32) >> bits) & 1).astype(cand_logits.dtype) # (bs,K,81,9) # Positive-weighted BCE. Candidate masks are sparse (~1-3 of 9 digits # "on"), so plain BCE collapses to predicting all-zeros. Up-weighting # the positive (candidate-present) term by pos_weight counteracts the # imbalance and forces the head to actually predict the candidate set. pos_weight = float(getattr(hyperparams, "aux_pos_weight", 5.0)) log_p = jax.nn.log_sigmoid(cand_logits) # log sigmoid(x) log_1mp = jax.nn.log_sigmoid(-cand_logits) # log(1 - sigmoid(x)) bce = -(pos_weight * cand_multi_hot * log_p + (1.0 - cand_multi_hot) * log_1mp) # (bs,K,81,9) # Only the empty cells are supervised: clue cells were sentinel-zeroed # in the data pipeline, so any cell whose target row is all-zero is a # clue and must not contribute. This makes the effective target # (#empty cells) x 9 per puzzle rather than 81 x 9. cell_mask = (cand_targets > 0).astype(cand_logits.dtype) # (bs, K, 81) # Delta weighting: consecutive stages are near-duplicates (at K=12 only # ~14% of cells change per stage, so ~97% of the target bits are copies # of the previous slot). Weighting the unchanged cells below 1 stops the # objective from being satisfied by echoing slot j-1. Slot 0 has no # predecessor, so it is fully weighted. delta_bg = float(getattr(hyperparams, "aux_delta_bg", 1.0)) if delta_bg != 1.0: changed = jnp.concatenate( [jnp.ones_like(cand_targets[:, :1], dtype=bool), cand_targets[:, 1:] != cand_targets[:, :-1]], axis=1) cell_w = cell_mask * (delta_bg + (1.0 - delta_bg) * changed.astype(cand_logits.dtype)) else: cell_w = cell_mask cell_denom = jnp.maximum(cell_w.sum(axis=2), 1e-6) # (bs, K) bce_per_cell = bce.mean(axis=3) # (bs, K, 81), mean over 9 digits bce_per_slot = ((bce_per_cell * cell_w).sum(axis=2) / cell_denom) # (bs, K), mean over empty cells if backtrack: # Strict stage-replay: supervise ONLY the stage-t readout = the last # active latent slot per example (index k-1), whose candidate target # is that stage's grid. Earlier slots and the LM CE are dropped, so # this step purely re-derives "apply f exactly t times -> stage-t # readout." Injection above still spans all active slots, so the # recurrence reaching slot k-1 is intact. k_per = active_full.sum(axis=1) # (bs,) last_idx = (k_per - 1)[:, None] # (bs, 1) slot_sel = ((jnp.arange(num_slots)[None, :] == last_idx) & active_full) # (bs, K) slot_active = slot_sel.astype(bce_per_slot.dtype) else: slot_active = active_full.astype(bce_per_slot.dtype) # (bs, K) aux_denom = jnp.maximum(slot_active.sum(), 1.0) avg_aux_loss = (bce_per_slot * slot_active).sum() / aux_denom # Backtrack (replay) steps train only from the readout (no LM CE). ce_term = 0.0 if backtrack else avg_ce_loss total_loss = ce_term + aux_weight * avg_aux_loss return total_loss, (pred_logits, avg_ce_loss, avg_aux_loss) # Compute the learning rate and perform gradient descent step = state.step lr = learning_rate_fn(step) (loss, aux), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) pred_logits, ce_loss, aux_loss = aux grads = jax.lax.pmean(grads, "batch") new_state = state.apply_gradients(grads=grads) # Update training metrics metrics = { "step": step, "loss": loss, "learning_rate": lr, "ce_loss": ce_loss, "aux_loss": aux_loss, "pred_logits": pred_logits, "weights": inputs.shape[0] } return new_state, metrics, pred_logits