Feat multihost - #239
Open
jurajHasik wants to merge 6 commits into
Open
Conversation
Root cause of multi-host failures: importing kfac-jax reverts jax to the legacy pmap implementation (jax_pmap_shmap_merge=False), which does not support multi-process runs in jax 0.9.1. Re-enable the new implementation in deepqmc.__init__ and adapt the host-side code to its stricter multi-process semantics (no eager ops on device-sharded arrays, no implicit resharding of committed arrays): - parallel: use public jax.distributed.initialize; correct local_device_ids (CUDA_VISIBLE_DEVICES ids are renumbered within the process); allow manual (non-SLURM) launches via DEEPQMC_NUM_PROCESSES/DEEPQMC_PROCESS_ID and JAX_COORDINATOR_ADDRESS - parallel: broadcast before tiling in replicate_on_devices(globally=True), route inputs through host numpy (committed inputs cannot feed pmap) - parallel: select_one_device fetches the addressable shard; new select_local_entries undoes all_gather on-device - fit: keep EWM data on the host, index inside pmap (also fixes global-vs-local device count in log_weight indexing), device_get observable samples for host-side consumers - train: pmapped equilibration criterion, host-side tau readout - log: avoid deepcopy of sharded arrays, deserialize tau via replicate_on_devices, clear error on batch/device-count mismatch - optimizer: pass momentum as jnp scalar (python scalars rejected by the new pmap in multi-process runs) Verified with 2 processes x 2 CPU devices (default ansatz and psiformer): training (KFAC), pretraining, equilibration, evaluation, checkpoint save/restore round trip; global stats identical across processes; single-host H2 energy converges to -1.174 Ha over 1000 steps. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
train_from_checkpoint walks the chain of restarts backwards to recover the original training config. A self- or mutually-referential task.restdir in a stored config (e.g. created by rolling-symlink restart patterns, or by relative restdir values resolved from a different working directory) made this loop spin forever. Track the visited directories and raise a clear error instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
On non-interactive runs (e.g. SLURM batch jobs) tqdm progress bars are disabled (disable=None means disable on non-TTY), and the only logged progress line was gated on the EWM error estimate halving, which occurs geometrically rarely. Long healthy runs therefore produced no progress output in the log file at all. Log a heartbeat every 1/5 of the pretraining steps and every 1/100 of the training steps. Also use the loop step number instead of pbar.n in the progress message: tqdm does not update n for disabled bars, so logged steps were stuck at 1 on non-interactive runs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Restarting from a run that ended before writing its first checkpoint (e.g. killed during pretraining or equilibration) failed with a bare IndexError. Explain what is missing and why instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add a CPU multi-host smoke test (two processes, one device each) covering training, checkpointing and restart from checkpoint. Co-Authored-By: Claude Opus <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Enable multi-host (pre)training leveraging jax.distributed