Index
base
¶
ShardingPolicy
¶
Bases: BaseModel
Execution parallelism; TP divides each data-parallel replica.
ShardingPlan(tp: Sequence[Rule] = (), zero: Sequence[Rule] | None = None)
dataclass
¶
Logical dimension mappings for tensor and optimizer-state parallelism.
By default ZeRO uses the TP dimension choices on the batch mesh axis. Unannotated tensors remain replicated regardless of the rules.
Topology
¶
Bases: BaseModel
new(chip: Chip, shard: ShardingPolicy | None = None) -> Topology
classmethod
¶
Create a Topology instance based on the current JAX device configuration.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
chip
|
Chip
|
The chip type being used. |
required |
shard
|
ShardingPolicy | None
|
Parallelism policy. The batch axis uses all remaining devices. |
None
|
TheoreticalFLOPS
¶
Bases: BaseModel
Theoretical FLOP/s, recorded in float TFLOP/s.
Dense peaks per physical chip: BF16 matrix operations with FP32 accumulation, and native FP32 arithmetic (not TF32). None means unknown or unpublished.
JobSpec
¶
Bases: BaseModel
user provided specification for a job
Node
¶
Bases: BaseModel
a stable identifier for a particular node
serialize() -> str
¶
serialize the node to a string
deserialize(s: str) -> Self
classmethod
¶
deserialize a node from a string
next() -> Node
¶
return the next node in the sequence.
You should NOT call this method unless you know what you are doing. The coherency of this method is important and if you tick accidentally it may break the DAG (e.g., either creating a node that's not stored or creating extra nodes.)
update(node: Node) -> None
¶
Update the live clock without replacing references held by consumers.
local(root_dir: str, work_dir: str) -> HardwareResult
¶
Detect hardware using JAX device information. Creates one ClusterMachine per host, with host index matching jax.process_index(). Assumes paths are identical across all hosts.