


Edit Distributed Machine Learning Training and Inference Code Without Breaking It
Use when changing software that trains or runs machine learning models across many processes or accelerators. It asks before risky changes, keeps model weight versions matched, keeps collective calls in the same order on every rank, protects config compatibility, and reports skipped tests.
To use this AI rule
## Read the Code First and Ask Before Changing Shared Config or APIs
Read the relevant implementation, neighboring patterns, configuration contracts, and tests before editing. Use the project environment manager, formatter, linter, and pre-commit stack; do not bypass hooks.
Never hardcode credentials, machine paths, service endpoints, or cluster topology. Do not guess scheduler settings or rebuild accelerator drivers and runtime stacks as a troubleshooting shortcut.
Obtain a decision before:
- changing shared configuration structures;
- adding dependencies;
- altering launcher or scheduler behavior;
- deleting, renaming, or incompatibly changing public APIs.
If required hardware or cluster access is unavailable, state which integration checks were skipped and why. Do not convert an unverified distributed assumption into code.
Route work by subsystem: training engine, inference engine, algorithm, launcher, scheduler, dataset, reward, workflow, or distributed debugging. Use the project’s specialized implementation guide or domain review path when one exists.
## Keep Each Subsystem Separate and Tag Data with Weight Versions
Keep orchestration, rollout generation, training, inference, infrastructure, model adaptation, rewards, datasets, and cross-cutting utilities as distinct responsibilities.
The trainer coordinates the loop, data loading, and workflows. Rollout workflows define asynchronous episode generation. Inference engines generate outputs and receive weight updates; training engines consume rollout tensors, compute updates, and publish new weight versions.
In asynchronous systems, attach explicit weight-version metadata to generated data and updates. Never train on or evaluate against silently mismatched model versions.
Emit metrics through the standard tracker, persist artifacts under the configured output root, and use the canonical save and recovery path for checkpoints. Observability and recovery are part of the execution contract.
Prefer composition and delegation over deep inheritance. Keep extension boundaries small enough that engines, workflows, and model adapters can be tested independently.
## Write Fast Accelerator Code That Avoids Device-to-Host Syncs
Follow existing naming, typing, and logging conventions in the owning module. Use explicit type hints and imports; never use wildcard imports. Keep heavyweight optional dependencies inside the functions or code paths that require them.
Use the project logger instead of raw printing. Include component and distributed-rank context in log records so concurrent output remains attributable.
In accelerator hot paths:
- avoid device-to-host synchronization such as extracting scalar or list values for logging;
- batch tensor operations instead of looping over individual elements in Python;
- set device and data type explicitly;
- assert expected tensor shapes at boundaries.
Keep rollout workflows non-blocking. Await asynchronous engine calls and file operations; do not introduce synchronous I/O into an asynchronous episode path.
For a new feature, follow the nearest working module pattern, add focused tests, and avoid unrelated refactoring.
## Change Config Fields and Public APIs Without Breaking Users
Order configuration fields from required, to common optional, to rare optional, then internal fields. Validate in initialization with a clear exception that names the invalid value or constraint.
Add public fields with safe defaults. Deprecate before removal, and avoid changing established field types. Public configuration classes must document constraints and accepted values.
Represent finite command-line choices with constrained types rather than free-form strings. Reuse canonical configuration dataclasses instead of creating parallel structures.
Treat shared configuration, launchers, schedulers, and public API names as high-impact surfaces. Make the compatibility effect explicit in the change description and obtain approval before destructive or incompatible edits.
When the correct behavior cannot be decided from current contracts, leave a clearly scoped follow-up marker and report the unresolved constraint instead of inventing a default.
## Keep Process Groups and Collective Calls Consistent Across Ranks
Never create a global process group at import time. Pass the intended process group explicitly and query rank against that group when group membership matters.
Keep device-mesh dimensions aligned with the declared parallelism model, including data sharding, tensor, context, expert, and expert-tensor axes where applicable.
Collective operations are global contracts:
- every participating rank must call the same all-reduce in the same order;
- select the reduction operation deliberately, such as sum versus mean;
- broadcasts must name an explicit source rank;
- barriers are diagnostic tools, not a routine synchronization strategy.
Diagnose hangs by checking mismatched collective sequences. Diagnose wrong results by checking reduction semantics and weight versions. Diagnose out-of-memory failures by checking sharding, tensor placement, and device ownership before reducing model functionality.
Use detailed distributed, communication-library, and synchronous-launch diagnostics only for targeted debugging, and remove or disable them after the issue is understood.
## Test Distributed Code and List Skipped Suites in the Pull Request
Name tests for behavior, condition, and expected outcome, and structure them as arrange, act, assert. Mark slow, asynchronous, accelerator, and distributed tests so default CI can select the correct suites.
Skip unavailable accelerator or multi-node tests gracefully with an explicit reason. Use the framework’s fake process group for unit-level distributed behavior; do not mock deep sharding or distributed-tensor internals whose semantics the test is meant to verify.
Compare tensors with explicit relative and absolute tolerances. Prefer managed temporary paths and scoped monkeypatching.
Run the narrow unit suite first, then the directly affected accelerator or distributed suite when hardware permits. Execute the full formatting and pre-commit stack before submission.
Generate command-line reference documentation through the canonical generator. Build release documentation through the project wrapper that performs required static setup and packaging, not a lower-level documentation command that omits project steps.
In the pull request, identify the issue, risk areas, exact test commands, skipped suites with reasons, and any hardware limitation. Use conventional, imperative commit subjects and remove work-in-progress commits before review.
