nemo_rl.data.energon.topology

Map policy parallel ranks to logical Energon loader copies.

Module Contents

Classes

LoaderCopy	One logical data shard and the policy ranks that consume it.
DataLoaderPlacementPlan	Stable mapping from logical data shards to policy ranks.
DataLoaderTopologyMapper	Build a logical loader placement for a policy sharding layout.
DefaultDataLoaderTopologyMapper	Create one loader per policy DP replica on TP0/PP0/CP0.

Functions

resolve_topology_mapper

Resolve the Stage 1 topology mapper.

API

class nemo_rl.data.energon.topology.LoaderCopy

One logical data shard and the policy ranks that consume it.

logical_rank: int

None

logical_world_size: int

None

owner_rank: int

None

delivery_ranks: tuple[int, ...]

None

class nemo_rl.data.energon.topology.DataLoaderPlacementPlan

Stable mapping from logical data shards to policy ranks.

copies: tuple[nemo_rl.data.energon.topology.LoaderCopy, ...]

None

placement_hash: str

None

property logical_world_size: int

copy_for_logical_rank(

logical_rank: int,

) → nemo_rl.data.energon.topology.LoaderCopy

Return the copy for logical_rank.

class nemo_rl.data.energon.topology.DataLoaderTopologyMapper

Bases: typing.Protocol

Build a logical loader placement for a policy sharding layout.

map(

sharding: nemo_rl.distributed.named_sharding.NamedSharding,

) → nemo_rl.data.energon.topology.DataLoaderPlacementPlan

class nemo_rl.data.energon.topology.DefaultDataLoaderTopologyMapper

Create one loader per policy DP replica on TP0/PP0/CP0.

map(

sharding: nemo_rl.distributed.named_sharding.NamedSharding,

) → nemo_rl.data.energon.topology.DataLoaderPlacementPlan

Map each DP coordinate to one owner and all replica consumers.

nemo_rl.data.energon.topology.resolve_topology_mapper(

name: str,

) → nemo_rl.data.energon.topology.DataLoaderTopologyMapper

Resolve the Stage 1 topology mapper.