Distributed API¶
Tools for multi-GPU training, elastic scaling, and gRPC-based actor workers.
Multi-GPU Training¶
rlox.distributed.multi_gpu.MultiGPUTrainer
¶
Wraps any rlox trainer for multi-GPU training via PyTorch DDP or FSDP.
Parameters¶
trainer_cls : type
Trainer class (e.g. PPOTrainer, SACTrainer).
env : str
Gymnasium environment ID.
config : dict, optional
Config overrides.
backend : str
PyTorch distributed backend ("nccl" for GPU, "gloo" for CPU).
strategy : str
Parallelism strategy: "ddp" or "fsdp" (default "ddp").
train(total_timesteps: int) -> dict[str, float]
¶
Run distributed training.
Only rank 0 returns full metrics; other ranks return reduced metrics.
Helpers¶
rlox.distributed.multi_gpu.is_main_rank() -> bool
¶
Return True if this process is rank 0 or distributed is not active.
Use this to guard logging, checkpointing, and evaluation that should only run once.
rlox.distributed.multi_gpu.reduce_metrics(metrics: dict[str, torch.Tensor], op: dist.ReduceOp = dist.ReduceOp.SUM) -> dict[str, torch.Tensor]
¶
rlox.distributed.multi_gpu.launch_elastic(trainer_fn: Callable[[], None], min_nodes: int = 1, max_nodes: int = 4, nproc_per_node: int = 1) -> None
¶
Launch fault-tolerant elastic training using torch.distributed.run.
This is a convenience wrapper around torch.distributed.launcher.api
for launching rlox trainers with elastic scaling support.
Parameters¶
trainer_fn : callable Zero-argument function that creates and runs a trainer. Will be invoked once per worker process. min_nodes : int Minimum number of nodes to start training (default 1). max_nodes : int Maximum number of nodes for elastic scaling (default 4). nproc_per_node : int Number of worker processes per node (typically num GPUs).
Raises¶
ValueError
If min_nodes > max_nodes.
RuntimeError
If the elastic launcher is not available.