code/train/trainer.py

# 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):
    """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). Gradients flow through all passes (full BPTT).

    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]
        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)

        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