← projects2024 · Research / RL

ClusterOrch-Gym

a reinforcement learning benchmark for training AI agents to diagnose and fix distributed GPU training failures.

distributed training fails sideways. one sick gpu slows a collective, a timeout trips on another node, the job retries, and by the time someone gets paged the first useful signal is buried under the cascade. clusterorch-gym gives an agent a repeatable place to learn that sequence.

the environment

the core is a discrete-event simulator that lines up nccl logs, kernel events, nvidia-smi samples, network counters, and checkpoint history on one clock. fault injectors can drop a link, poison a gradient, exhaust memory, slow a rank, or surface a gpu xid while the rest of the job keeps moving.

an observation is a rolling window over device utilization, memory pressure, collective latency, link health, rank progress, and checkpoint age. the window matters because an isolated timeout says very little; the order in which several weak signals arrive is usually the diagnosis.

actions include isolating a node, moving work to a spare, rolling back to a checkpoint, changing collective topology, restarting a rank, or waiting for another sample. every action has a cost in lost tokens, stale work, or healthy capacity, so an agent that panics at every spike learns an expensive lesson.

the reward follows useful training time, recovery latency, checkpoint loss, and false interventions. the environment keeps diagnosis and recovery separate in the trace, which makes it possible to see whether a policy understood the fault or merely found a restart button that worked.

training

the reference agent uses ppo with a temporal policy over the telemetry window. the causal mask limits it to signals already available at that step, and parallel simulator instances let the same fault arrive with different timing, load, and checkpoint freshness.

evaluation replays held-out scenarios across single-node and multi-node layouts. the score records recovered training time, interventions, and the amount of healthy work discarded along the way.

failure modes

The benchmark includes 15 distinct failure modes organized into 4 categories:

hardware: gpu xid errors, pcie link degradation, ecc pressure, and thermal throttling.

network: nccl watchdog timeouts, infiniband link flaps, incast congestion, and routing failures.

software: cuda out-of-memory, version mismatches, stuck collectives, and checkpoint i/o hangs.

silent: nan injection, stragglers, and corrupted state that keeps the job alive long enough to waste more compute.

github.com/Pavitrakus/clusterorch-gym ↗ · who else is here