mojolearn.LanguageModelTrainer#

mojolearn.LanguageModelTrainer#

Methods

attention_stage_report()

Memory held by the open session's attention stages.

close()

Copy the device state to the host and release the device.

evaluate(ids)

Mean loss on ids without changing any state.

export_checkpoint(path)

Write the complete state to path as a JSON checkpoint.

export_checkpoint_binary(path)

Write the complete state to path as a binary checkpoint.

export_gradients([named])

The pre-update gradient of the last completed step of a resident session.

export_state()

The complete training state as fresh copies.

from_checkpoint(path, *[, resident])

Restore a trainer from a JSON checkpoint written by export_checkpoint().

from_checkpoint_binary(path, *[, resident])

Restore a trainer from export_checkpoint_binary().

from_checkpoint_bytes(encoded, *[, resident])

Restore a trainer from the bytes of a JSON checkpoint.

load_state_dict(state)

Replace the complete training state.

logits(ids, *[, lengths])

Float32 logits for every position, without changing any state.

next_bytes(ids, *[, lengths])

Greedy next token after each row.

parameter_registry([shape])

Name, shape, offset and size of every parameter tensor.

run_metadata()

Profile, configuration, schedule and binary details of this run.

save_checkpoint(path)

Same as export_checkpoint().

state_dict()

The complete training state; same as export_state().

train_step(ids)

Run one training step: mean cross-entropy loss and an AdamW update.

Attributes

data_schedule

A copy of the data_schedule descriptor, as checkpoints keep it.

parameters_

A copy of the current flat float32 parameters.

step_

Number of completed training steps.

alias of SmallByteLanguageModelTrainer