kempnerforge.training.runtime

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

Functions

pp_group(runtime)

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

setup_distributed(config)

Initialize the process group + device mesh, then validate the config against it.

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.setup_distributed(config)[source]

Initialize the process group + device mesh, then validate the config against it.

Parameters:

config (JobConfig)

Return type:

RuntimeContext