kempnerforge.training.runtime

Process-level context for a training job: ranks, device, mesh, pipeline.

Functions

apply_matmul_precision(train)

Run fp32 matmuls on TF32 tensor cores when train.tf32 is set; otherwise do nothing.

pp_group(runtime)

The PP process group for this rank, derived on demand.

setup_distributed(config)

Validate the config for the world size, then initialize the process group + device mesh.

Classes

PipelineBundle

Pipeline-parallel state.

RuntimeContext

Where this process sits: ranks, device, and the mesh built for them.
class kempnerforge.training.runtime.RuntimeContext[source]

Bases: object

Where this process sits: ranks, device, and the mesh built for them.

rank: int
local_rank: int
world_size: int
device: torch.device
device_mesh: DeviceMesh | None
__init__(rank, local_rank, world_size, device, device_mesh)
Parameters:
Return type:

None

class kempnerforge.training.runtime.PipelineBundle[source]

Bases: object

Pipeline-parallel state. Built only when distributed.pp > 1.

The PP process group is not held here: it is derived from the mesh at the call sites that need it, so a job that never runs a step never slices it.

rank: int
size: int
schedule: Any
__init__(rank, size, schedule)
Parameters:
Return type:

None

kempnerforge.training.runtime.pp_group(runtime)[source]

The PP process group for this rank, derived on demand.

Parameters:

runtime (RuntimeContext)

Return type:

Any

kempnerforge.training.runtime.apply_matmul_precision(train)[source]

Run fp32 matmuls on TF32 tensor cores when train.tf32 is set; otherwise do nothing.

Parameters:

train (TrainConfig)

Return type:

None

kempnerforge.training.runtime.setup_distributed(config)[source]

Validate the config for the world size, then initialize the process group + device mesh.

Parameters:

config (JobConfig)

Return type:

RuntimeContext