dew.objectives.lm

Name	Summary
`TEXT_KEY`	Batch key the token pipeline packs `[B, seq_len + 1]` int32 ids under.
`IndexerTraining`	Train DeepSeek-V3.2’s lightning indexer, one stage at a time.
`LMObjective`	Train a next-token model: shifted cross entropy, teacher-forced scoring, optional previews.
`LMRunConfig`	A run, plus the language model’s own knobs.
`Perplexity`	Report exp of the cross entropy per counted target over a whole pass.
`Samples`	Configure the text preview drawn once per event.
`perplexity`
`prompt_batch`	Build `[B, P]` int32 ids from one prompt, or several of the same length.

TEXT_KEY

attribute source

TEXT_KEY = 'text'

Batch key the token pipeline packs [B, seq_len + 1] int32 ids under.

IndexerTraining

dataclass source

class IndexerTraining(phase: Literal['warmup', 'sparse'], weight: float = 1.0)

Train DeepSeek-V3.2’s lightning indexer, one stage at a time.

The two stages of the continued pre-training (arXiv 2512.02556, section 2.1.1). warmup trains a fresh indexer alone: the model runs dense attention with the indexer scoring beside it (an mla mixer with the indexer’s heads and no top-k), every other weight is frozen, and the loss is the KL of the indexer’s softmax from the dense attention distribution over every allowed key. sparse trains everything: the model selects its top-k (a mixer with index_topk), the cross entropy trains the main weights, and the KL over the selected keys alone trains the indexer, whose inputs are detached so neither reaches the other. weight scales the KL term (MaxText’s indexer_loss_scaling_factor); the reference sets the indexer’s pace by its learning rate, 1e-3 for the 1000 warm-up steps and 7.3e-6 for the sparse stage.

LMObjective

class source

class LMObjective(
    model,
    seq_len: int,
    *,
    ema_decay: float | None = 0.999,
    pad_id: int | None = None,
    head_chunks: int = 4,
    samples: Samples | None = None,
    pretrained: Variables | None = None,
    balance_rate: float | None = None,
    aux_loss_alpha: float | None = None,
    seq_aux: bool = True,
    loss_role: Role | None = None,
    mtp_weight: float | None = None,
    z_loss: float = 0.0,
    router_z_loss: float = 0.0,
    qk_stats: bool = False,
    indexer: IndexerTraining | None = None,
    trainable: PathFilter | None = None,
    token_accuracy: bool = True,
)

Train a next-token model: shifted cross entropy, teacher-forced scoring, optional previews.

LMObjective.held_variables

def held_variables() -> Variables | None

Return the checkpoint a continued-pretraining run starts from.

Bound as the initializer’s argument this reaches the trainer’s state JIT as data; read off self inside a nullary trace it would be compiled into the executable as a constant.

LMObjective.init

def init(key, variables: Variables | None = None) -> Variables

LMObjective.policy

def policy(params: Variables, sampling: Sampling = Sampling()) -> TextGeneration

Expose the model over this training tree as a generation task.

A rollout binds one snapshot of the policy and draws every completion from it; the result records the actual and raw-policy likelihoods the objective’s ratio needs.

LMObjective.pipeline

def pipeline(
    state: TrainState,
    *,
    ema: bool = True,
    processor: Processor | None = None,
) -> TextGeneration

Publish the decoder over the state’s weights as a generation task.

It samples and is budgeted the way this objective’s previews are, and processor decodes.

LMObjective.token_scores

def token_scores(
    params,
    tokens,
    train: bool = False,
    rngs=None,
    segment_ids=None,
    positions=None,
    routing: bool = False,
    depths: bool = False,
    roles=None,
    qk_stats: bool = False,
    indexer: bool = False,
    layers: Sequence[int] = (),
    routes: tuple[jax.Array, jax.Array | None] | None = None,
)

Score per-token next-token cross entropy over a [B, seq_len + 1] batch.

Returns Scores: the losses, the weight of each target, whether each prediction was right, the states behind them, and what routing, depths, qk_stats, indexer and layers asked for. layers names the model’s layers whose output states to keep, layers_N in its tree, as a feature distillation reads them.

A packed batch carries segment_ids for the same rows. The last token of a document does not predict the first of the next one, so that transition is dropped from the loss and the accuracy, and the model reads the per-document positions for its rotary angles. A chat batch carries roles for the same rows; with loss_role set, only the targets whose role matches keep their weight.

routes replays a rollout engine’s expert choices: [B, seq_len + 1, layers, top_k] ids aligned with tokens, and [B, seq_len + 1] booleans marking the ids the record covers (None for all). Every router selects those experts instead of its own top-k and still weights them from its scores (dew.nn.moe.Routes); the stack slices the record by layer however it runs.

LMObjective.per_token_log_probs

def per_token_log_probs(
    params: Variables,
    tokens: jax.Array | ModelInputs,
    *,
    left_padding: jax.Array | None = None,
) -> jax.Array

Score raw policy likelihoods aligned to next-token targets.

Explicit left-padding counts move real context to position zero before scoring. Returned slots whose input is padding are zero and unscored.

LMObjective.sampled_log_probs

def sampled_log_probs(
    params: Variables,
    scores: Scores,
    tokens: jax.Array,
    support: tuple[jax.Array, jax.Array] | None = None,
    temperature: float = 1.0,
) -> jax.Array

Each next-token target’s likelihood as the sampler that drew it saw it.

scores is token_scores over tokens, [B, S + 1]; the result is [B, S]. At unit temperature without support that is the raw policy, -scores.losses. temperature divides the capped logits (head_logits). support is the per-row ragged (ids, columns) pair sessions.pack builds, [B, C] each, columns the column in tokens of the id each kept id belongs to; a target with entries is renormalized over them (support_log_probs).

LMObjective.loss

def loss(params, batch, step: Step) -> tuple[Mean | LMStatistics, Aux[Variables]]

LMObjective.predict

def predict(
    params,
    batch,
    step: Step,
    *,
    train: bool,
    layers: Sequence[int] = (),
) -> tuple[Mean, Aux[Variables], Prediction]

Score the loss with the logits, the target weights and the outputs of layers behind it, for a teacher to compare (Objective.predict).

The logits are the whole [B, seq_len, vocab] fp32 tensor the chunked loss never holds; a distillation’s KL reads every column. aux_loss_alpha’s router terms carry their own normalisation, so they cannot ride a distillation’s token mass and are refused; balance_rate balances without a loss term. The indexer warm-up scores no token, so it has nothing to distil.

LMObjective.reduce_loss

def reduce_loss(stats: Mean | LMStatistics) -> tuple[jax.Array, jax.Array]

LMObjective.apply_effects

def apply_effects(variables: Variables, effects: Variables) -> Variables

LMObjective.evaluate

def evaluate(params, batch, step: Step)

Score the complete batch teacher-forced, using EMA when present.

LMObjective.preview

def preview(params, batch, step: Step, *, scored=None)

Sample the configured prompt once, then decode only on process zero.

An objective whose EMA holds a frozen reference draws from the live policy instead, which is what _ema_is_reference says.

LMRunConfig

dataclass source

class LMRunConfig(
    model: ModelConfig = (lambda: ModelConfig('causal_transformer'))(),
    data: DataSpec = TokenWindows(),
    optim: OptimConfig = (lambda: OptimConfig(learning_rate=0.0006, weight_decay=0.1, clip_grads=1.0))(),
    trainer: TrainerConfig = TrainerConfig(),
    objective: str = 'lm',
    lora: LoRA | None = None,
    tokenizer: str = 'byte',
    ema_decay: float | None = 0.999,
    sample_prompt: str = '',
    sample_tokens: int = 128,
    sampling: Sampling = (lambda: Sampling(temperature=0.8, top_k=40))(),
    pretrained: str | None = None,
    balance_rate: float | None = None,
    aux_loss_alpha: float | None = None,
    seq_aux: bool = True,
    router_z_loss: float = 0.0,
    mtp_weight: float | None = None,
    indexer: IndexerTraining | None = None,
    token_accuracy: bool = True,
    block_prompt_tokens: int = 256,
    block_canvas_size: int | None = None,
)

A run, plus the language model’s own knobs.

objective: str: Loss convention: lm, masked_diffusion (MDLM), or block_diffusion (the official DiffusionGemma fine-tuning objective).
tokenizer: str: What the ids were written with: ‘byte’, or an HF tokenizer name.
ema_decay: float | None: None disables EMA; 1.0 retains a frozen copy.
sample_prompt: str: Prompt the validation samples continue; empty continues a newline.
sample_tokens: int: Tokens generated per validation sample; 0 logs no text.
sampling: Sampling: The preview policy, recorded with the run for inference.
pretrained: str | None: Hugging Face decoder to continue training: a hub repo id, repo@revision (a branch, tag or commit), or a local directory in that layout. A run records a hub repo as repo@commit, the commit it resolved to. The checkpoint decides the architecture, so —model.config may then carry max_seq_len alone.
balance_rate: float | None: How far a sparse run moves each router’s balancing bias against its load every step (DeepSeek’s aux-loss-free balancing). Needs a mixture with bias=True; unset leaves the bias where it is.
aux_loss_alpha: float | None: The expert balance loss’s weight (LMObjective.aux_loss_alpha); with —no-seq-aux it is the Switch loss over the step’s routed positions, lm-engine’s router_aux_loss_coef. Unset adds no balance loss.
seq_aux: bool: Form the balance loss within each sequence (DeepSeek V2) rather than over the whole step.
router_z_loss: float: The routers’ z-loss weight (LMObjective.router_z_loss); lm-engine uses 0.1 times its aux coefficient. Zero adds nothing.
mtp_weight: float | None: DeepSeek’s lambda on the multi-token-prediction term. Needs a model with num_nextn_predict_layers above zero; unset leaves the term out.
indexer: IndexerTraining | None: DeepSeek-V3.2’s lightning-indexer phase: indexer:indexer-training --indexer.phase warmup freezes everything but the indexer of a model whose mla mixer names the indexer’s heads and no top-k; sparse trains the whole model on its top-k with the KL beside the cross entropy. Unset trains no indexer term.
token_accuracy: bool: Report the argmax accuracy beside the loss; False skips the argmax over every logit it costs.
block_prompt_tokens: int: Clean prompt prefix in a block-diffusion token row.
block_canvas_size: int | None: Training canvas width; None uses the checkpoint canvas length.

Perplexity

class source

class Perplexity

Report exp of the cross entropy per counted target over a whole pass.

Every batch weighs by its own count of counted targets, so a packed or padded pass whose batches differ in size is scored per token, and a batch with no counted target contributes nothing.

Perplexity.merge

def merge(
    accumulated: tuple[float, float],
    contribution: tuple[float, float],
) -> tuple[float, float]

Perplexity.finalize

def finalize(accumulated: tuple[float, float]) -> float

Samples

dataclass source

class Samples(
    prompt: Sequence[int] | Sequence[Sequence[int]],
    max_new_tokens: int,
    sampling: Sampling = Sampling(),
    decode: Callable[[list[int]], str] = lambda ids: str(ids),
)

Configure the text preview drawn once per event.

Prompts contain token IDs, with equal lengths for multiple prompts. This display count does not limit the teacher-forced scoring population.

perplexity

function source

def perplexity() -> Perplexity

prompt_batch

function source

def prompt_batch(prompt) -> jax.Array

Build [B, P] int32 ids from one prompt, or several of the same length.