TrainState#

class TrainState(schema_version: int = 1, epoch: int = 0, iter_in_epoch: int = 0, global_step: int = 0, samples_seen: int = 0, world_size: int = 1, batch_size_per_rank: int | None = None, grad_accum_step: int = 0, best_metric_value: float | None = None, rng_state: dict[str, ~typing.Any] | None = None, sampler_state: dict[str, ~typing.Any] | None = None, dataloader_state: dict[str, ~typing.Any] | None = None, extra_state: dict[str, ~typing.Any] = <factory>)[source]#

Bases: object

Serializable resume point for trainer, sampler, dataloader, and RNG state.

classmethod from_state_dict(state_dict: dict[str, Any]) → TrainState[source]#: Create TrainState from checkpoint data.

classmethod from_trainer(trainer: Any) → TrainState[source]#: Capture the next resume position from a live trainer.

load_state_dict(state_dict: dict[str, Any]) → None[source]#: Load state with compatibility for older iteration key names.

state_dict() → dict[str, Any][source]#: Return a checkpoint-friendly dictionary representation.

batch_size_per_rank: int | None = None#

best_metric_value: float | None = None#

dataloader_state: dict[str, Any] | None = None#

epoch: int = 0#

extra_state: dict[str, Any]#

global_step: int = 0#

grad_accum_step: int = 0#

iter_in_epoch: int = 0#

property iteration: int#: Backward-compatible alias for iter_in_epoch.

rng_state: dict[str, Any] | None = None#

sampler_state: dict[str, Any] | None = None#

samples_seen: int = 0#

schema_version: int = 1#

world_size: int = 1#