Add opt-in adaptive data-parallel search trials - #868
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0b8d6e175c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if world_size % per_node: | ||
| raise ValueError( | ||
| f"world_size={world_size} cannot be placed on homogeneous " | ||
| f"{per_node}-GPU nodes" | ||
| ) |
There was a problem hiding this comment.
On allocations where the per-node GPU count does not divide the planned world size, this aborts even though the trial fits. For example, the default allowed sizes select world_size=8 for one survivor in a two-node, six-GPU-per-node allocation, but eight processes could run as four per node; this branch instead raises because 8 % 6 != 0. Account for node topology while planning, or launch a uniform subset of each node's GPUs rather than requiring every participating node to be full.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
🤖 Addressed in cd88672. The Slurm allocation executor now plans only uniformly placeable world sizes and computes an explicit (nodes, processes_per_node) placement. On a 2×6-GPU allocation, world size 8 now launches as 2 nodes × 4 ranks instead of requiring divisibility by all six GPUs per node. Impossible custom world sizes are filtered during planning, with a warning and fallback to the largest placeable compatible size.
0b8d6e1 to
4d3ce68
Compare
|
🤖 A real W&B observability failure surfaced in the Torch run: |
4d3ce68 to
c90c8ba
Compare
c90c8ba to
cd88672
Compare
|
🤖 Addressed the W&B observability failure reported in the experiment. The fix landed in the stacked base PR (#865) as bf9a1e2 and is inherited here: the potentially long search run ID is no longer added as a W&B tag, while the complete ID remains available through the W&B group and structured |
8b49218 to
7a84d2a
Compare
f088b69 to
652d674
Compare
7a84d2a to
be7ed07
Compare
|
🤖 A final retry-safety review found and fixed a possible local deadlock in be7ed07. If a persisted adaptive plan requires more GPUs than a retry exposes, the local device pool now fails immediately with recovery guidance instead of waiting forever. The docs also explain how to derive |
|
@codex will you review this PR? |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: be7ed070e8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ).stdout.splitlines() | ||
| if not hosts: | ||
| raise RuntimeError(f"Slurm returned no hosts for {node_list!r}") | ||
| node_rank = int(os.environ["SLURM_NODEID"]) |
There was a problem hiding this comment.
When concurrent srun --exclusive steps receive a subset of a larger allocation, SLURM_NODEID is the node's job-allocation index rather than a dense rank within this step. A two-node trial placed on allocation nodes 2 and 3 therefore passes node ranks 2 and 3 to torchrun --nnodes=2, so rendezvous cannot complete. Use the step's dense SLURM_PROCID, as the existing launchers do in scripts/slurm_apptainer_train.sbatch:461 and scripts/empireai_nvl72_train.sbatch:224.
Useful? React with 👍 / 👎.
| if adaptive_batching: | ||
| train_batch_sampler = FixedGlobalBatchSampler.from_datasets( |
There was a problem hiding this comment.
For candidates containing BatchNorm—including the shipped Samudra configurations—enabling this adaptive path does not preserve the claimed training trajectory. At world_size=1 the model uses ordinary per-local-batch statistics, while distributed rungs convert it to SyncBatchNorm in Trainer.__init__, whose statistics cover local_batch_size * world_size; gradient accumulation and the fixed sampler cannot make those forward-pass statistics equivalent. Consequently, expanding a survivor changes its optimization semantics despite holding the effective global batch fixed; either reject such models for adaptive scaling or keep normalization behavior invariant.
Useful? React with 👍 / 👎.
| complete = len(indices) // self.global_batch_size | ||
| optimizer_batches.extend( | ||
| indices[start : start + self.global_batch_size] | ||
| for start in range( | ||
| 0, | ||
| complete * self.global_batch_size, |
There was a problem hiding this comment.
When training has multiple resolution groups, computing complete independently for every group discards up to effective_global_batch_size - 1 samples per group, whereas the previous sampler allowed homogeneous microbatches from different groups to participate in one accumulation cycle. For example, two 48-sample groups with a global batch of 64 previously provide three optimizer batches but this sampler provides only two; if every group is smaller than 64, construction fails even though their combined microbatches are sufficient. This changes both epoch sample coverage and optimizer-step counts, so preserve group-homogeneous microsteps while carrying their tails across accumulation cycles.
AGENTS.md reference: AGENTS.md:L18-L20
Useful? React with 👍 / 👎.
|
🤖 Final experiment-integration review at |
|
🤖 Empirical follow-up: the exact |
1 participant
Summary
adaptive_data_parallelsearch resource policytorchrunjobs on local GPUs and across homogeneous nodes in an existing Slurm allocationScientific semantics
The planner computes
gradient_accumulation_steps = effective_global_batch_size / (batch_size * world_size). A fixed-global-batch sampler partitions the same epoch-seeded optimizer batches across ranks and accumulation microsteps, so changing world size does not change the samples or optimizer-step count. Learning-rate and scheduler settings are therefore left unchanged. DDP synchronization is skipped on non-final accumulation microsteps withno_sync().Candidate configs must use
backend: auto, which lets the existinginit_distributed_mode()machinery select single-process execution or initialize fromtorchrunenvironment variables. Fixed anchors remain single-GPU. The separately submitted Slurm-array executor rejects this policy because those jobs do not share an allocation.Batch-size and placement warnings
If
effective_global_batch_sizeis not divisible bybatch_size * requested_world_size, the controller emits a warning with the candidate and configured batch size, preserves that choice, and selects the largest smaller compatible world size. If the batch size cannot realize the target even on one GPU, adaptive scaling is disabled for that candidate and its original accumulation setting is retained.The allocation executor separately constrains plans to world sizes that can be distributed uniformly without exceeding the visible GPUs on any node. If topology rules out an otherwise compatible world size, the controller warns and selects the largest placeable alternative.
Testing
uv run python -m pytest tests/test_search_resources.py tests/test_samplers.py tests/test_search.py -q— 82 passeduv run python -m pytest -m "not manual and not cuda" -n auto— 480 passed, 2 skipped, 10 xfaileduvx pre-commit run --all-files— all hooks passedTests cover 1/2/4-rank optimizer-batch invariance, adaptive expansion, compatible and irreconcilable batch warnings, retry plan persistence, local multi-GPU launch construction, partially occupied multi-node placement, and Slurm/torchrun rendezvous construction. This machine has no CUDA/Slurm allocation, so the real 8- and 16-GPU smoke runs remain deployment checks.
Stack
This is intentionally based on #865 and contains only the adaptive data-parallel follow-up. It can be retargeted to
mainafter #865 merges.