Exp docs 0.2.0 preview - #146
Open
laserkelvin wants to merge 45 commits into
Open
Conversation
All CUDA dependencies (cuml-cuXX, cuequivariance-ops-torch-cuXX, pylibraft-cuXX, nvidia-dali, cupy, nvidia-physicsnemo) now publish real wheels to public PyPI, so the extra NVIDIA index is no longer required to resolve them. Remove the `nvidia` uv index, correct the now-stale `index-strategy` comment, and re-lock — uv.lock now sources solely from PyPI and the PyTorch CUDA indexes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
nvidia-physicsnemo installs and imports on macOS (pure py3-none-any wheel; its CUDA-only deps are gated behind cuXX extras), and it is already an unconditional base dependency. Drop the CPU-only `_install_physicsnemo_docs_stub` so autodoc renders against the real package instead of fake signatures. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the inline-script install builder with an accessible selector matrix (package manager / accelerator / optional extras) whose logic lives in a deferred static JS file and whose styling moves to custom.css. The generated commands now emit the PyTorch CUDA index for pip installs (`--extra-index-url https://download.pytorch.org/whl/cuXXX`) so CUDA extras actually resolve, while uv relies on `--torch-backend`. Drop the redundant pypi.nvidia.com index from the prose examples now that those packages ship on public PyPI. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ProfilerHook was removed; the safety-and-monitoring example failed to import and broke the docs gallery build. Switch to the drop-in StageTimingHook (same constructor, summary(), and CSV schema) in the example script, and update the README prose that still referenced the old hook. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sphinx autosummary writes per-module .rst into docs/modules/**/generated and docs/modules/**/_generated at build time. Ignore these so build artifacts are not committed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The AtomicDataZarrWriter/AtomicDataZarrReader docstrings introduced the store-layout ASCII tree with a single colon and an indented block, which RST parses as a definition list — collapsing the tree's indentation and monospacing in the generated API pages. Wrap the tree in an explicit ``.. code-block:: text`` directive so it renders as a verbatim monospace literal block (the ``text`` lexer also avoids the project's default Python highlighting being applied to the tree characters). Note: because this is a docstring-only change, autodoc's incremental build may not pick it up — rebuild clean (``make docs-clean docs``) to see it. Fixes NVIDIA#98 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add autodoc-pydantic to the docs group and route pydantic models through it: the autosummary class template emits `autopydantic_model` for models with no leftover Parameters/Attributes section, a field-type override renders the clean jaxtyping annotation (not the raw Annotated[...] with Field/PlainSerializer noise), and validators/config/JSON are collapsed. Also dedent docstrings in the dataclass-table directive so dataclass Attributes descriptions render. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move per-field descriptions into Field(description=...) so autodoc-pydantic renders them, and remove the now-redundant numpy Parameters/Attributes docstring sections. Expand the class docstrings (AtomicData, ModelConfig, NeighborConfig, TrainingStrategy, ValidationConfig, OptimizerConfig, BufferConfig, the Zarr configs, CLI job specs, and training hooks) with purpose/usage prose, Examples, and Notes; add a "See also" from the Zarr configs to the compression tuning guide. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
De-duplicate hook API pages: dynamics/api.rst no longer regenerates the dynamics hooks (references dynamics-hooks instead), and the shared StageTimingHook/TorchProfilerHook are documented only in hooks-api, removing the cross-page duplicate-object warnings. Add autosummary entries so NeighborConfig and BufferConfig get rendered pages (BufferConfig's cross-reference path is corrected to nvalchemi.dynamics.base). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Convert the two reporting-flow HTML diagrams to top-down .. graphviz:: diagrams matching every other userguide/module page (shared green/dark theme), and drop the now-unused .reporting-flow CSS. Also fix the broken "Designing Rich layouts" cross-reference by adding an explicit label. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Document what each loss computes with a shared notation (B graphs indexed by i, N_i atoms indexed by a, hats for predictions, valid-entry sums, the Huber function): per-atom vs graph-balanced energy MSE/MAE/Huber, force MSE/Huber and L2-norm, and stress Frobenius terms. Add equations for the CosineWeight (half-cosine ramp) and PiecewiseWeight (step function) schedules under the same schedule-index notation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Explain the checkpoint-saving logic: a spec-based (non-pickle) capture of model, optimizer/scheduler, training counters, and CheckpointableHook state, the intentionally-excluded training_fn/loss_target_assembler, and the manifest indexing behavior. Add a labeled cross-reference to the training guide's checkpoint-semantics section. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Document the profiler lifecycle (stages bound, start/step/finalize), output/rank-subdir layout, and registration, with Examples and Notes — matching the usage docstrings added to the other pydantic hooks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Document how the data classes relate and differ from their PyTorch counterparts. Dataset and DataLoader now explain the Reader -> Dataset (-> MultiDataset) -> DataLoader chain, AtomicData/Batch semantics, threaded CUDA-stream fused prefetch vs worker forks, and how samplers plug in. MultiDataset documents the global index space and field-name contract, and the two multi-dataset samplers contrast per-sample stochastic mixing (MultiDatasetSampler as sampler=) with fixed per-batch composition (MultiDatasetBatchSampler as batch_sampler=), each with examples and reciprocal See Also links. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Lift InMemoryDataset's class docstring to match the sibling data-pipeline classes: position it as the memory-resident counterpart to Dataset in the Reader -> Dataset/InMemoryDataset (-> MultiDataset) -> DataLoader pipeline, note that it is interchangeable via BatchDatasetProtocol, and give when-to-use guidance versus Dataset. Add a reciprocal cross-reference from Dataset to InMemoryDataset. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Convert unicode math symbols in docstrings to LaTeX (``:math:``) where they render, or clean ASCII in the non-rendered private modules; finish aligning parameter names, shapes, defaults, and described behavior with the code across data, models, dynamics, hooks, training, and distributed; and fix an invalid inline-literal in the ``current_dd_context`` docstring. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Add an Integrator Algorithms page documenting each integrator/optimizer with its scheme, equations, units, and the backing ``nvalchemiops`` kernels in collapsible panels. Add distributed API references split into the spatial domain-decomposition building API and the general-purpose runtime utilities (``DistributedManager`` + rank/world/device resolvers). Expand the models module page with per-model physics, units, install notes, and kernel panels. Wire all of the above into the module toctrees. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Convert RST colon-roles to MyST ``{py:...}`` roles, disambiguate the install
links, repair reference-label anchors, correct the NPT pressure units
(eV/A^3) and the DemoDynamics naming, align the dynamics/training hook API
pages with the current signatures, and make graphviz diagram text readable on
the dark theme.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Give graphviz nodes with light fills an explicit dark font color so their labels are readable on the dark documentation theme, and lengthen a too-short section-title underline in the 2-D FIRE/NVT distributed example. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Contributor
|
Too many files changed for review. ( Bypass the limit by tagging |
Fix I001 import-block ordering in docs/conf.py and docs/sphinxext.py and ruff-format normalization in docs/sphinxext.py and nvalchemi/training/strategy.py, surfaced by the all-files ruff hooks in the CI lint job. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
… types) A mechanical signature-vs-docstring audit surfaced docstrings that claimed values the code does not use. Correct the parameter defaults that are actually None (NeighborListHook.stage, LevelSchema group/dtype maps, BaseLevelStorage.device, configure_dataloader.shuffle) to state the real default and what it resolves to; fix the vec9 pressure-tensor shape ([M, 9] not [M, 3, 3]) in the NPT/NPH barostat kernels; correct DemoModel.forward's positions type and force/energy shapes and document compute_forces; document hybrid_forces on the Ewald/PME wrappers; and clarify ReportingOrchestrator.error_policy's enum default. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
ys-teh
reviewed
Jul 23, 2026
The ``*Compile`` tests build ``torch.compile(fn, backend="cudagraphs")`` functions on CUDA. The cudagraph tree manager retains each graph's output tensors across invocations, so under a ``pytest-testmon`` selective run a prior compile test's retained output is still live in the shared cudagraph pool when the next test captures a graph, tripping ``check_memory_pool``'s "Detected N tensor(s) in the cudagraph pool not tracked as outputs" check. The full suite happens to order/GC around it, so it fails only in selective (PR) runs. Add an autouse fixture that resets dynamo and frees the cudagraph pool before and after each compile test, making any test subset/order hermetic. Cost is paid only by the few compile tests; verified harmless on the CPU compile path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Restructure the class docstring into a scannable three-case list (what you set decides the behaviour) so the freeze/trainable precedence is clear at a glance, and tighten the freeze_patterns/trainable_patterns field descriptions to match. Addresses ys-teh's PR review comment. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
ys-teh
reviewed
Jul 25, 2026
ys-teh
left a comment
ys-teh
left a comment
Collaborator
There was a problem hiding this comment.
Skimmed through a few pages, and everything looks fine to me so far. Additional visual would be helpful in general, though some of the graphs and flowcharts may just not have rendered correctly on my end. I also left some minor comments around.
Collaborator
There was a problem hiding this comment.
minimal end-to-end usage example?
Comment on lines
653
to
664
Collaborator
There was a problem hiding this comment.
This looks the same as the above block code?
Collaborator
There was a problem hiding this comment.
Fix the numbering here and the subsequent ones.
ys-teh
requested changes
Jul 27, 2026
ys-teh
reviewed
Jul 27, 2026
Review feedback from ys-teh: * ``CheckpointHook`` takes ``checkpoint_dir``, not ``root_folder`` — fix the TrainingStrategy docstring example (``root_folder`` remains correct for ``save_checkpoint``/``load_checkpoint``, which is why the two diverged). * Drop three duplicated rows from the agent-skills table. * Drop a duplicated ``SPEC_*`` import block in the design overview, and fix the section numbering: the UMA section was inserted as "4b" to avoid renumbering and was missing from the contents table. It is now section 5, later sections shift to 6/7/8 (with their subsections), all in-page section cross-references are updated, and the contents table lists it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Review feedback from ys-teh:
* Add rendered previews of the two built-in Rich dashboards (training and
dynamics) to the reporting guide, so readers can see what a reporter looks
like before investing in one. The SVGs are produced by the documented
``RichReporter.preview()`` call itself, so they can be regenerated.
* Replace the remaining ``{eval-rst}``-wrapped ``.. graphviz::`` blocks with
MyST-native ``{graphviz}`` fences, matching the 18 diagrams that already used
the native form. These blocks render identically in the built docs; the
wrapper only made the source harder to read in review (GitHub does not render
Sphinx directives, which is why they appeared as raw text).
* Cross-link the distributed guide's overview list to the sections it
describes, including the minimal end-to-end example.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
The first fix keyed the cudagraph reset off a *Compile class or a test name containing "compile", which cut the failures from 13 to 1 but left TestUtilsCompile::test_scatter_reduce_sum_compiles[cuda] failing: the tensor left in the shared cudagraph pool comes from test/distributed, which pytest collects *before* test/dynamics and whose compiling tests match neither predicate. Key the reset off the test module instead — any module whose source builds a torch.compile-d callable (28 modules) is bracketed — and fix the reset order so in-flight work is drained and unreferenced tensors collected *before* the cudagraph trees are torn down; a dangling reference otherwise pins the pool and the teardown cannot free it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Rich's SVG export leaves a trailing space on one line per file, which the pre-commit trailing-whitespace hook rewrites and CI lint therefore rejects. The stripped files are still valid XML with identical rendered content (no <text> element is affected). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Bracketing every module that calls torch.compile cleared the cudagraph failures (0 remaining) but pushed the GPU test job past its 45-minute limit: 803 tests paid two GC + empty_cache cycles each. Only the cudagraphs backend and mode="reduce-overhead" route through cudagraph_trees; the eager / aot_eager compiles in the distributed suite never allocate from the cudagraph pool, so they cannot leak into it and do not need bracketing. Scoping to modules that can build a pool drops the bracketed set from 803 tests to 361, and dropping empty_cache() removes the dominant per-reset cost — reset_cudagraph_trees() is what actually releases the pool, which the previous run confirmed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
…lures Bracketing all 361 tests in the cudagraph-using modules kept the pool clean (0 errors) but blew the 45-minute GPU job limit: torch._dynamo.reset() discards the compile cache, so every bracketed test recompiles (~1.7s each) and the suite reached only 84%% in 42 minutes against a 29m31s baseline. Restrict the reset to the ~16 tests that actually build cudagraph-backed compiles — they recompile anyway, so the reset is nearly free there — and stop suppressing reset failures. The traceback shows the surviving failure happens in CUDAWarmupNode.run -> check_memory_pool, i.e. the pool was already dirty on entry, which is what a silently-failed teardown reset would produce. The warning names the test so the next run identifies it instead of leaving the downstream test to report a confusing error. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
The surviving failure reports only the address of a tensor left live in the cudagraph pool. check_memory_pool already calls gc.collect() and torch.cuda.synchronize() before it errors, so the tensor is genuinely referenced rather than uncollected garbage, and adding further resets only masks it. Enable cudagraph_trees_history_recording so the failure names the capture that allocated the tensor. Only cudagraph captures allocate from that pool, so this identifies the leaking test directly. The flag scopes allocator history to cudagraph record/warmup, so it does not slow the wider suite. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
…orkaround
Regenerating the lockfile on this branch silently upgraded warp-lang from
1.14.0 to 1.15.0. That upgrade, not the documentation, caused the cudagraph
failures: the CI runs differ in exactly one package.
PR NVIDIA#150 (passes): torch 2.13.0+cu130, triton 3.7.1, warp-lang 1.14.0
this branch (fails): torch 2.13.0+cu130, triton 3.7.1, warp-lang 1.15.0
The failing test compiles scatter_reduce_per_graph, whose _segmented_sum
is a custom_op wrapping wp.from_torch and a Warp kernel launch. Under
1.15.0 that path leaves an allocation live in the cudagraph memory pool which
is not tracked as a graph output, so the next capture fails
check_memory_pool. warp-lang >= 1.13.0 has no upper bound, so only the
lockfile pinned it.
Restore the lockfile from 0.2.0-rc and re-lock, which adds only
autodoc-pydantic (plus its pydantic-settings dependency) for the docs
build and leaves every other pin untouched. Revert the conftest reset fixture
added earlier: it was masking this dependency change rather than fixing a test
isolation bug.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Per review feedback, add a user guide for the (de)serialization machinery. It is a cross-cutting feature rather than a training-only one — the same spec layer backs model wrappers, optimizers, schedulers, loss terms and weight schedules, hooks, and the training CLI — so the mechanics are documented once in their own page and referenced from training. The new page covers why recipes are used instead of pickle, the anatomy of a ``BaseSpec`` (``cls_path`` plus one field per constructor kwarg), which value types survive a round trip, the two-step precedence for deriving a model spec (``checkpoint_spec()`` first, attribute introspection second) and the save-time rebuild check that guards the fallback, how to register custom types, what is deliberately never serialized, and the design rules that follow from all of it. Training gains a Reproducibility section stating the five conditions a run must meet to be resumable, and noting that ``Omitting model spec`` warnings should be treated as errors. Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
Collaborator
Author
|
@ys-teh unfortunately it seems like line numbers aren't available for the comments you made and I left unresolved... |
Collaborator
May I check which comment this refers to? |
ys-teh
approved these changes
Jul 28, 2026
Brings in the ghost-width validation fix (NVIDIA#152) and the CI supply-chain hardening (pinned action SHA, sha256 installer checksum). Conflict in nvalchemi/distributed/config.py was additive on both sides: upstream appended a sentence to the ghost_width docstring entry while this branch added strategy/compile/require_nondegenerate entries at the same insertion point. Both retained. Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
15 tasks
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
2 participants
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.
ALCHEMI Toolkit Pull Request
Description
This PR revamps the documentation for the 0.2.0 release, including user-facing and API documentation across the board
Type of Change
Related Issues
Changes Made
Testing
make pytest)make lint)Checklist
Additional Notes
Tip
This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.