fix(API): resolve torch_musa._initialized dynamically so checkpoint preserves device RNGFix/initialized stale reexport

[yinjiew]

Fixes #150.

Problem

torch_musa/__init__.py:120 re-exports _initialized from torch_musa.core._lazy_init
with from ... import. That binds a snapshot of the value False, not the name, so the
_initialized = True that _lazy_init() performs at core/_lazy_init.py:137 can never
reach torch_musa.__dict__["_initialized"].

torch.musa is this module (torch.__setattr__("musa", sys.modules[__name__]),
__init__.py:57), and torch.utils.checkpoint decides whether to save and restore the
device RNG state with getattr(device_module, "_initialized", False) — in both
CheckpointFunction.forward and _checkpoint_without_reentrant_generator.

So on MUSA, torch.utils.checkpoint(..., preserve_rng_state=True) silently degrades to
CPU-RNG-only. Every stochastic op inside a checkpointed region (nn.Dropout, stochastic
depth, Gumbel-softmax, any custom noisy activation) draws different random numbers in the
backward recompute than it did in the forward pass, producing a wrong gradient with no
error, no exception and no warning. Measured on 2x MTT S5000 with torch_musa 2.7.1 at
ResNet-18 scale, max abs gradient difference against the uncheckpointed reference: 2.6
for a stochastic activation, 139.7 with a learnable branch. Training still converges to
something, which is why this can go unnoticed indefinitely.

is_initialized() is a function, so it re-reads the live global and returns the correct
answer. That asymmetry is why every internal consumer behaves correctly and only upstream
PyTorch, reading the raw attribute, is affected. torch.cuda does not have this bug because
its _initialized and _lazy_init() live in the same module.

Affected in every release I could inspect — 2.7.1, 2.9.1.post1 and 2.11.0.post1 (this
branch). core/_lazy_init.py is byte-identical across all three, and PyTorch v2.7.1, v2.9.1
and v2.11.0 all gate on the same attribute. Nothing in torch_patches/ touches
torch/utils/checkpoint.py.

Fix

Stop re-exporting the flag; forward it with a module-level __getattr__ (PEP 562) so
core/_lazy_init remains the single source of truth. That is correct automatically for
every path, including re-initialization after fork, because there is only ever one binding.

Deleting the re-export is required, not cosmetic: __getattr__ is consulted only when
normal attribute lookup fails, so a leftover stale entry in torch_musa.__dict__ would keep
shadowing the forwarder.

The other five names this file took from core/_lazy_init (_lazy_init, _lazy_call,
_is_in_bad_fork, is_initialized, musart) are now read off the same _lazy_init_mod
handle, so there is one consistent way this file reaches into that module. Binding those
five statically is safe — each is a callable assigned once at import time and never
rebound, and rebinding a name to a function object cannot go stale. _initialized is
precisely the exception, being a plain mutable False/True, which is why it alone needs
the forwarder.

The alternative — writing torch_musa._initialized = True from inside _lazy_init() — is a
smaller diff, but it creates two bindings that must then be kept in sync at every site that
touches the flag, including bad-fork reset and any future teardown path. Happy to switch if
you prefer it.

torch_musa/core/mudnn.py already uses the sibling technique (a PropModule replacing
sys.modules[__name__]) so that torch.backends.mudnn.allow_tf32 reads through to C++
instead of returning a stale Python copy — same problem, same shape of solution.

Commits

Six commits, in three pairs — the original change and then the follow-up that addresses
your review of it. Happy to squash any or all of them if you would rather have a flatter
history.

commits	net effect
eb51bc3 + d452548	the fix: drop the _initialized re-export, add the __getattr__ forwarder, and route the remaining five names through the same _lazy_init_mod handle
a8dceb8 + 7c78f53	fix the _queue_calls typo in _lazy_init()'s global statement to _queued_calls, matching upstream torch/cuda/__init__.py. The declaration is a no-op for that name either way — it is only ever mutated with .append(), never rebound — so this is readability only, no behaviour change
d54b0c0 + ecb6793	regression tests

Tests

Neither existing test could have caught this:

• tests/unittest/amp/test_amp_checkpoint.py checkpoints a fully deterministic nn.Linear
  and asserts only inputs.dtype == output.dtype; test_amp_checkpoint_fp16 has no
  assertion at all.
• tests/unittest/miscs/test_musa_converter.py:70 does contain
  assert torch.musa._initialized, but inside the python_dst string literal — it is the
  expected output of the CUDA-to-MUSA source converter, compared as text and never
  executed.

Added:

• tests/unittest/standalone/test_lazy_init.py::TestLazyInit::test_lazy_init_flag — the
  existing test, extended to also assert torch_musa._initialized and
  torch.musa._initialized once initialization has happened. It lives in the directory
  where scripts/run_unittest.sh gives each file its own process, which is what makes an
  assertion about lazy-init state meaningful.

• tests/unittest/core/test_checkpoint.py (new), two tests:

  • test_checkpoint_preserves_device_rng, parametrized over both use_reentrant values —
    records what each pass drew from inside the checkpointed function and asserts the
    backward recompute replayed the forward's draws bitwise. This is the direct probe of
    the defect.
  • test_checkpoint_dropout_gradient_matches_eager — the black-box counterpart. An
    nn.Dropout mask cannot be read out of the module the way the draws list reads it
    out of a hand-rolled function, so gradient equality against an independently seeded
    eager run is the only way to verify replay for a real nn.Module.

  Picked up by the existing pytest ... tests/unittest/core line in
  scripts/run_unittest.sh, so no CI change is needed.

What I could and could not verify

I no longer have access to the MTT S5000 machine, so I have not executed these tests
against 2.11.0.post1 — please run them in CI. To be explicit about where every claim in
this PR comes from:

claim	how it was established
wrong gradients under checkpoint on MUSA	measured on hardware — 2x MTT S5000, torch_musa 2.7.1+5ee0a64 (reproduction and output in #150)
torch_musa._initialized stale while is_initialized() is correct	measured on hardware, same session
refreshing the flag restores bitwise noise replay	measured on hardware, same session
the same defect is present on this branch	source review — the re-export at __init__.py:120 is unchanged from 2.7.1
the Python mechanism itself	reproduced with no GPU and no PyTorch — self-contained script and output in #150
the branch applies to upstream/main	verified — merges clean, no conflicts
formatting	verified — black 24.2.0, the version pinned in .pre-commit-config.yaml, reports all four files unchanged. Re-checked after each round of review changes
pylint	could not run the hook locally (it needs torch/torch_musa importable). I checked the new code against tools/lint/pylintrc by hand: function-rgx / variable-rgx / argument-rgx / method-rgx, good-names, docstring-min-length, and the docparams accept-no-*-doc defaults

Nothing above is extrapolated; every claim is labelled with how it was obtained.

Behaviour change worth a release note

torch/utils/checkpoint.py also raises RuntimeError("PyTorch's device state was initialized in the forward pass of a Checkpoint, which is not allowed.") when
_initialized is truthy but the forward did not stash device state. With the stale flag that
guard could never fire; it now can. This matches CUDA and is upstream's intended behaviour,
and it is hard to reach in practice — the checkpointed region would have to be the first
thing in the process to touch a MUSA tensor — but it is a deliberate difference rather than a
pure bug fix.

Backports

The same re-export exists in the 2.9.1.post1 and 2.7.1 trees and this change applies there
too, but the repository has no release branches, so I can only target main. Happy to open
follow-up PRs if you ever cut them.

Environment

torch        2.7.1a0+gite2d141d  (measured)  /  consumer reviewed on v2.7.1, v2.9.1, v2.11.0
torch_musa   2.7.1+5ee0a64 (measured); 2.9.1.post1 and 2.11.0.post1 (source-reviewed)
GPU          2 x Moore Threads MTT S5000 (80 GB)
Python       3.10, Ubuntu container