nemo_automodel.components.distributed.ddp

Module Contents

Classes

Name	Description
`DDPManager`	Deprecated compatibility wrapper for PyTorch DDP.

Functions

Name	Description
`_resolve_ddp_device`	Validate distributed state and return the device owned by this rank.
`parallelize_ddp`	Apply activation checkpointing and PyTorch DDP from a typed config.

Data

logger

API

class nemo_automodel.components.distributed.ddp.DDPManager(
    config: nemo_automodel.components.distributed.config.DDPConfig
)

Deprecated compatibility wrapper for PyTorch DDP.

.. deprecated:: 0.7 Pass DDPConfig through the config-driven infrastructure and provide model-specific behavior with ModelParallelizer. This compatibility class is scheduled for removal in 0.8.

This manager wraps models with DistributedDataParallel for data-parallel distributed training.

Parameters:

config

DDPConfig

Configuration for DDP distributed training.

nemo_automodel.components.distributed.ddp.DDPManager.__getattr__(
    name
)

nemo_automodel.components.distributed.ddp.DDPManager._setup_distributed()

Initialize device configuration for DDP.

Sets the rank, world_size, and device based on the process group backend.

nemo_automodel.components.distributed.ddp.DDPManager.parallelize(
    model: torch.nn.Module,
    reapply_trainability: collections.abc.Callable[[torch.nn.Module], None] | None = None
) -> torch.nn.Module

Wraps the given model with DistributedDataParallel (DDP).

Moves the model to the initialized device before wrapping. For CUDA devices, the device id is passed to DDP as device_ids; for CPU, no device ids are provided.

Parameters:

model

torch.nn.Module

The PyTorch model to be wrapped.

reapply_trainability

Callable[[torch.nn.Module], None] | NoneDefaults to None

Optional callback that re-resolves parameter trainability after model surgery and before DDP construction.

Returns: torch.nn.Module

torch.nn.parallel.DistributedDataParallel: The DDP-wrapped model.

nemo_automodel.components.distributed.ddp._resolve_ddp_device() -> torch.device

Validate distributed state and return the device owned by this rank.

nemo_automodel.components.distributed.ddp.parallelize_ddp(
    model: torch.nn.Module,
    config: nemo_automodel.components.distributed.config.DDPConfig,
    device: torch.device | None = None,
    reapply_trainability: collections.abc.Callable[[torch.nn.Module], None] | None = None
) -> torch.nn.Module

Apply activation checkpointing and PyTorch DDP from a typed config.

nemo_automodel.components.distributed.ddp.logger = logging.getLogger(__name__)