Fine-tuning¶

laya.train is imported on its own (from laya.train import finetune, TrainConfig); it is not re-exported from the top-level package. laya-train is its command line. The fine-tuning guide shows both in use.

Running a fine-tune¶

TrainConfig `dataclass` ¶

TrainConfig(
    epochs: int = 4,
    micro_batch: int = 8,
    grad_accum: int = 8,
    encoder_lr: float = 2.5e-05,
    head_lr: float = 0.0001,
    min_lr: float = 1e-06,
    weight_decay: float = 0.01,
    grad_clip: float = 1.0,
    loss: str = "rlcd",
    rl_samples: int = 4,
    sigma_start: float = 0.4,
    sigma_end: float = 0.1,
    w_sph: float = 0.75,
    w_rps: float = 1.0,
    shuffle_options: Tuple[str, ...] = (),
    option_layout: Optional[str] = None,
    freeze_encoder: bool = False,
    calib_max: int = 400,
    calib_frac: float = 0.1,
    calib_seed: int = 20260922,
    seed: int = 0,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    amp: Optional[bool] = None,
    gradient_checkpointing: Optional[bool] = None,
    log_every: int = 100,
    label_smoothing: float = 0.0,
    text_column: str = "text",
    label_column: str = "label",
    question_id: str = "label",
    instructions: Optional[str] = None,
    eval_data: Optional[str] = None,
    target_error: float = 0.1,
    min_abstain_n: int = 10,
)

Training knobs. The defaults train the way the published notebook does.

The notebook's effective batch of 64 is 8 per micro-batch x 2 GPUs x 4 accumulation steps; on one device that is micro_batch=8, grad_accum=8.

max_len / head_max_len default to the checkpoint's own config. Whatever is used is written into the saved config, so inference sees the same budgets training did. amp defaults to on for CUDA only; gradient_checkpointing defaults to following amp. option_layout defaults to the checkpoint's own ("sequential" for every published one); "parallel" trains on common.parallel_layout and is written into the saved config.

finetune ¶

finetune(
    data: str,
    model_dir: str,
    output_dir: str,
    config: Optional[TrainConfig] = None,
    device: Optional[str] = "auto",
) -> Dict[str, Any]

Preprocess, train, calibrate and save; returns a summary of the run.

data is a JSONL file of {state, questions, gold} or {state, questions, expected} rows, or a CSV file with text and label columns. model_dir is a local checkpoint directory -- the layout laya.load reads. The result in output_dir loads with laya.load(output_dir).

dry_run ¶

dry_run(
    data: str, model_dir: str, config: Optional[TrainConfig] = None
) -> Dict[str, Any]

Inspect dataset and build items without loading model weights; returns dataset stats.

The two halves of `finetune`¶

train_model ¶

train_model(
    model,
    tok,
    items: Sequence[Dict[str, Any]],
    config: TrainConfig,
    device: device,
    max_len: int,
    head_max_len: int,
    on_epoch_end: Optional[Callable[[int, float], None]] = None,
    parallel: bool = False,
) -> List[float]

Train model in place on items; returns the mean loss of each epoch.

calibration_records ¶

calibration_records(
    model,
    tok,
    items: Sequence[Dict[str, Any]],
    device: device,
    max_len: int,
    head_max_len: int,
    batch_size: int = 16,
    parallel: bool = False,
)

(qtype, logits, target, k) records for fit_temperature_map, options in canonical order.

calibration_report ¶

calibration_report(
    records: Sequence[Tuple[int, Any, Any, int]],
    temperature: Sequence[float],
) -> Dict[str, Dict[str, Any]]

Per question type: how many calibration items there were, what was fitted, and what to doubt.

records are the calibration_records the temperatures were fitted on, and temperature the per-type scalars fit_temperature_map returned. Each type gets items, temperature, accuracy and mean_confidence (top probability at the fitted temperature) on the calibration items, and a list of issues, empty when nothing looks wrong:

fewer than MIN_TYPE_N items: the temperature was not fitted and stays 1.0;
fewer than CALIB_WARN_N items: it was fitted, on little evidence;
it landed on the [TEMP_MIN, TEMP_MAX] clamp;
it came back unchanged at 1.0 although it was fitted, which usually means the items gave the fit nothing to correct -- what a model already certain and right on them produces.

A type with no calibration items at all has no issues; it was not in the data.

Reading data¶

read_data ¶

read_data(
    path: str,
    text_column: str = "text",
    label_column: str = "label",
    question_id: str = "label",
    instructions: Optional[str] = None,
) -> List[Dict[str, Any]]

Read training/eval rows from either a JSONL file or a CSV file.

rows_from_csv ¶

rows_from_csv(
    path: str,
    text_column: str = "text",
    label_column: str = "label",
    question_id: str = "label",
    instructions: Optional[str] = None,
) -> List[Dict[str, Any]]

Read a CSV file with one labelled column per row into {state, questions, expected} rows.

Discovers distinct labels in label_column (preserving order of appearance) and builds a choice question with those labels as a list of criteria to render each option once. Uses utf-8-sig encoding so UTF-8 BOM headers from Excel are handled transparently.

Building training items¶

items_from_rows ¶

items_from_rows(
    tok,
    rows: Iterable[Dict[str, Any]],
    max_len: int,
    head_max_len: int,
    label_smoothing: float = 0.0,
) -> Tuple[List[Dict[str, Any]], Dict[str, int]]

Items from {state, questions, gold} or {state, questions, expected} rows; returns (items, skipped).

skipped counts, by reason, questions that were labelled but could not become an item: empty_text, empty_label, invalid_question and invalid_target (a ValueError from to_internal, target_from_gold or target_from_expected), plus the reasons make_item gives. Questions with neither a gold nor an expected entry are not counted -- the row simply does not label them.

target_from_gold ¶

target_from_gold(
    q: Dict[str, Any], gold_q: Dict[str, Any]
) -> List[float]

The normalised target distribution for an internal question (to_internal), in option order.

gold_q["probabilities"] is keyed by choice label, by "false"/"true" for noul, and by the level index as a string for score -- the {state, questions, gold} schema in docs/finetune.md. A choice label that is not a string is also looked up as its JSON key. A row whose probabilities sum to zero becomes uniform, as in the scripts. A negative or non-finite probability raises ValueError.

target_from_expected ¶

target_from_expected(
    q: Dict[str, Any], expected_val: Any, label_smoothing: float = 0.0
) -> List[float]

The normalised target distribution from a single expected answer, in option order.

expected_val is: - a choice label (str or value matching a criterion) for choice questions - a boolean for noul questions - an integer level index or level string for score questions

When label_smoothing > 0, smooths probability mass uniformly: target[i] = (1 - eps) * one_hot[i] + eps / K.

to_internal ¶

to_internal(qid: Any, question: Dict[str, Any]) -> Dict[str, Any]

The question as inference sees it: Agent._check_question then Agent._to_internal.

Training on any other rendering would teach the model option texts it never sees at inference -- a list-form choice, custom noul labels, a non-string instruction.

make_item ¶

make_item(
    tok,
    q: Dict[str, Any],
    target: Sequence[float],
    state_ids: List[int],
    head_max_len: int,
    max_len: Optional[int] = None,
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]

(item, None), or (None, reason) when the question cannot be trained on as written.

q is an internal question (to_internal). The reasons are target_mismatch (the target does not have one entry per option), options_collapsed (the head budget left two options with the same token span, #538) and, when max_len is given, options_beyond_max_len (the head is longer than max_len, so build_sequence drops the markers of the last options -- the case Agent refuses at inference with "only N of its M option markers fit"). Training on any of them would teach the model to tell apart options it cannot see as different, or crash the batch on a target longer than its markers.

encode_item ¶

encode_item(
    tok,
    item: Dict[str, Any],
    max_len: int,
    head_max_len: int,
    option_order: Optional[List[int]] = None,
    parallel: bool = False,
) -> Dict[str, Any]

The model input for an item, with options in option_order and the target to match.

build_sequence puts option option_order[s] in slot s, so the slot-ordered target is target[option_order[s]] -- the inverse of what unpermute_probs does at inference.

draw_option_order ¶

draw_option_order(
    item: Dict[str, Any], rng: Random, shuffle_types: Sequence[str]
) -> Optional[List[int]]

A random permutation for items whose type is in shuffle_types, else None (canonical).

split_calibration ¶

split_calibration(
    items: Sequence[Any], calib_max: int, calib_frac: float, seed: int
) -> Tuple[List[Any], List[Any]]

(train, calibration), with the calibration slice taken before any training.

Fitting temperatures on items the run has trained on measures the fit, not the calibration (docs/finetune.md). The split depends only on seed and the item count, so every DDP rank and every rerun hold out the same items.

Losses¶

soft_ce_loss ¶

soft_ce_loss(
    logits: Tensor, target: Tensor, mask: Tensor
) -> torch.Tensor

Cross-entropy against a target distribution, over each row's own options only.

rlcd_loss ¶

rlcd_loss(
    logits: Tensor,
    target: Tensor,
    mask: Tensor,
    qtype: Tensor,
    sigma: float,
    samples: int = 4,
    w_sph: float = 0.75,
    w_rps: float = 1.0,
) -> torch.Tensor

The notebook objective: a policy-gradient term over noisy logits plus soft cross-entropy.

Zero-mean Gaussian noise (scaled by sigma, projected onto each row's options) is added to detached logits; each sample is rewarded by proper_reward against the target, and the normalised advantage weights the Gaussian log-density of the sample.

Checkpoints¶

resolve_checkpoint_dir ¶

resolve_checkpoint_dir(
    model_id_or_path: str, token: Optional[str] = None
) -> str

Resolve a local path, model alias ('multilingual', 'english') or Hub repo ID to a local directory.

load_checkpoint ¶

load_checkpoint(model_dir: str)

(model, tokenizer, cfg) from a local checkpoint directory or Hub ID, model on CPU in fp32.

save_checkpoint ¶

save_checkpoint(model, tok, cfg: Dict[str, Any], path: str) -> None

Write the layout laya.load reads: weights in fp16, encoder config, tokenizer, config.