Class ParallelState

java.lang.Object

smile.llm.parallel.ParallelState

public final class ParallelState
extends Object

Per-rank view of a ParallelConfig mesh.

Phase 1: globalRank == tpRank with ppRank == 0. Phase 2 multi-node PP will set globalRank = ppRank * tpSize + tpRank.

Constructor Summary

Constructors

Constructor	Description
ParallelState(ParallelConfig config, int tpRank)	Phase-1 constructor (ppRank=0, dpRank=0).
ParallelState(ParallelConfig config, int tpRank, int ppRank, int dpRank)	Full mesh-coordinate constructor.

Method Summary

Modifier and Type	Method	Description
static void	clearCurrent()	Clears the thread-local parallel state.
ParallelConfig	config()	Returns the process-mesh configuration.
static ParallelState	current()	Returns the state bound to the calling thread, if any.
byte	deviceIndex()	Returns the CUDA device index for this rank.
int	dpRank()	Returns this rank's data-parallel index.
int	globalRank()	Returns the global rank in the process mesh.
boolean	isFirstStage()	Returns whether this rank is the first pipeline stage.
boolean	isLastStage()	Returns whether this rank is the last pipeline stage.
boolean	isTpRoot()	Returns whether this rank is the TP root (tpRank == 0).
int	ppRank()	Returns this rank's pipeline-parallel index.
static void	setCurrent(ParallelState state)	Binds this state to the calling thread for collective helpers.
int	tpRank()	Returns this rank's tensor-parallel index.
int	tpSize()	Returns the tensor-parallel size.

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

ParallelState

public ParallelState(ParallelConfig config,
 int tpRank)

Phase-1 constructor (ppRank=0, dpRank=0).

Parameters:

config - process-mesh configuration.

tpRank - tensor-parallel rank.

ParallelState

public ParallelState(ParallelConfig config,
 int tpRank,
 int ppRank,
 int dpRank)

Full mesh-coordinate constructor.

Parameters:

config - process-mesh configuration.

tpRank - tensor-parallel rank.

ppRank - pipeline-parallel rank.

dpRank - data-parallel rank.

Method Details

config

public ParallelConfig config()

Returns the process-mesh configuration.

Returns:

parallel config.

tpSize

public int tpSize()

Returns the tensor-parallel size.

Returns:

TP size.

tpRank

public int tpRank()

Returns this rank's tensor-parallel index.

Returns:

TP rank.

ppRank

public int ppRank()

Returns this rank's pipeline-parallel index.

Returns:

PP rank.

dpRank

public int dpRank()

Returns this rank's data-parallel index.

Returns:

DP rank.

globalRank

public int globalRank()

Returns the global rank in the process mesh.

Returns:

global rank.

deviceIndex

public byte deviceIndex()

Returns the CUDA device index for this rank.

Returns:

device index.

isTpRoot

public boolean isTpRoot()

Returns whether this rank is the TP root (tpRank == 0).

Returns:

true if TP root.

isFirstStage

public boolean isFirstStage()

Returns whether this rank is the first pipeline stage.

Returns:

true if ppRank == 0.

isLastStage

public boolean isLastStage()

Returns whether this rank is the last pipeline stage.

Returns:

true if last PP stage.

setCurrent

public static void setCurrent(ParallelState state)

Binds this state to the calling thread for collective helpers.

Parameters:

state - state to bind, or null to clear.

current

public static ParallelState current()

Returns the state bound to the calling thread, if any.

Returns:

current state, or null if unset.

clearCurrent

public static void clearCurrent()

Clears the thread-local parallel state.