JoinHook

class tensorplay.distributed.algorithms.JoinHook[source]

This defines a join hook, which provides two entry points in the join context manager.

Entry points : a main hook, which is called repeatedly while there exists a non-joined process, and a post-hook, which is called once all processes have joined.

To implement a join hook for the generic join context manager, define a class that inherits from JoinHook and override main_hook() and post_hook() as appropriate.

main_hook() → None[source]

Call this hook while there exists a non-joined process to shadow collective communications in a training iteration.

Training iteration i.e., in one forward pass, backward pass, and optimizer step.

post_hook(is_last_joiner: bool) → None[source]

Call hook after all processes have joined.

It is passed an additional bool argument is_last_joiner, which indicates if the rank is one of the last to join.

Parameters:

is_last_joiner (bool) – True if the rank is one of the last to join; False otherwise.