# utils.distributed – Axolotl

# utils.distributed

`utils.distributed`

Utilities for distributed functionality.

## Functions

Name | Description
--- | ---
barrier | Acts as a barrier to wait for all processes. This ensures that all processes
cleanup_distributed | Destroy process group if torch distributed is initialized. Called in training early
compute_and_broadcast | Compute a value using the function ‘fn’ only on the specified rank (default is 0).
gather_from_all_ranks | Run a callable ‘fn’ on all ranks and gather the results on the specified rank.
gather_scalar_from_all_ranks | Run a callable ‘fn’ on all ranks and gather the results on the specified rank.
is_distributed | Check if distributed training is initialized.
is_main_process | Check if the current process is the main process. If not in distributed mode,
reduce_and_broadcast | Run a callable ‘fn1’ on all ranks, gather the results, reduce them using ‘fn2’,
zero_first | runs the wrapped context so that rank 0 runs first before other ranks

### barrier

    utils.distributed.barrier()

Acts as a barrier to wait for all processes. This ensures that all processes reach the barrier before proceeding further.

### cleanup_distributed

    utils.distributed.cleanup_distributed()

Destroy process group if torch distributed is initialized. Called in training early termination or when training successfully completes.

### compute_and_broadcast

    utils.distributed.compute_and_broadcast(fn)

Compute a value using the function ‘fn’ only on the specified rank (default is 0). The value is then broadcasted to all other ranks.

Args: - fn (callable): A function that computes the value. This should not have any side effects. - rank (int, optional): The rank that computes the value. Default is 0.

Returns: - The computed value (int or float).

### gather_from_all_ranks

    utils.distributed.gather_from_all_ranks(fn, world_size=1)

Run a callable ‘fn’ on all ranks and gather the results on the specified rank.

Args: - fn (callable): A function that computes the value. This should not have any side effects. - rank (int, optional): The rank that gathers the values. Default is 0. - world_size (int, optional): Total number of processes in the current distributed setup.

Returns: - A list of computed values from all ranks if on the gathering rank, otherwise None.

### gather_scalar_from_all_ranks

    utils.distributed.gather_scalar_from_all_ranks(fn, world_size=1)

Run a callable ‘fn’ on all ranks and gather the results on the specified rank.

Args: - fn (callable): A function that computes the value. This should not have any side effects. - rank (int, optional): The rank that gathers the values. Default is 0. - world_size (int, optional): Total number of processes in the current distributed setup.

Returns: - A list of computed values from all ranks if on the gathering rank, otherwise None.

### is_distributed

    utils.distributed.is_distributed()

Check if distributed training is initialized.

### is_main_process

    utils.distributed.is_main_process()

Check if the current process is the main process. If not in distributed mode, always return `True`.

We use a simpler logic when the distributed state is not initialized: we just log on the 0-th local rank.

#### Returns

Name | Type | Description
--- | --- | ---
| bool | `True` if the current process is the main process, `False` otherwise.

### reduce_and_broadcast

    utils.distributed.reduce_and_broadcast(fn1, fn2)

Run a callable ‘fn1’ on all ranks, gather the results, reduce them using ‘fn2’, and then broadcast the reduced result to all ranks.

Args: - fn1 (callable): A function that computes the value on each rank. - fn2 (callable): A reduction function that takes a list of values and returns a single value. - world_size (int, optional): Total number of processes in the current distributed setup.

Returns: - The reduced and broadcasted value.

### zero_first

    utils.distributed.zero_first(is_main)

runs the wrapped context so that rank 0 runs first before other ranks