fix: guard tokenizer save with rank 0 to avoid NFS deadlock

[dafu-wu]

What does this PR do ?

Adds a rank-0 guard around tokenizer saving so that only one process writes the tokenizer directory during checkpointing.

Every rank previously called tokenizer.save_pretrained() on the same tokenizer_path. Unlike model/optimizer state, the tokenizer is replicated rather than sharded, and save_pretrained writes rank-independent filenames (tokenizer_config.json, ...), so all ranks raced on the same files.

save_pretrained is not atomic: it opens with O_TRUNC, writes, and may read files back. Concurrent writers therefore contend on the inode lock. On a hard-mounted NFS share this deadlocks: one rank blocks in do_truncate -> nfs_setattr holding the inode write lock while the others block in nfs_start_io_read -> down_read, all in uninterruptible disk sleep (D state). Those ranks never return, so the driver waits forever in ray.get() and the whole job hangs at checkpoint time with no error and no timeout.

Observed failure

On an 8-node / 32-rank DTensor v2 SFT run (gemma-4-31B-it), training went silent right after Saving checkpoint for step 500... for 35+ hours. 31 ranks had finished and returned to Ray's main_loop; 4 ranks on a single node were stuck:

 PID  STAT WCHAN                    ELAPSED      COMMAND
1533  DNl  nfs_start_io_read      1-22:05:56  ray::DTensorPolicyWorkerV2.save_checkpoint
1535  DNl  rpc_wait_bit_killable  1-22:05:56  ray::DTensorPolicyWorkerV2.save_checkpoint

Kernel stacks (/proc/PID/stack) show the writer/reader interlock, all four on tokenizer_config.json:

# writer
rpc_wait_bit_killable [sunrpc] -> nfs4_do_setattr [nfsv4] -> nfs_setattr
  -> notify_change -> do_truncate -> handle_truncate -> do_open

# readers
down_read -> nfs_start_io_read [nfs] -> nfs_file_read -> vfs_read

dmesg confirms: rwsem_down_read_slowpath, blocked for more than 245 seconds.

How the guard went missing

nemo_automodel already guards these artifacts on rank 0 in ConsolidatedHFAddon.pre_save. But AutomodelCheckpointManager.save_checkpoint passes tokenizer=tokenizer if tokenizer_path is None else None to save_model(), and callers (e.g. sft.py) always pass tokenizer_path. So the addon receives tokenizer=None, its rank-0 branch never runs, and control falls through to a locally implemented, unguarded save_pretrained().

Scope of the guard

Deliberately tokenizer-only. Model and optimizer saves go through dcp.save, which is collective and already writes rank-disjoint files (shard-000NN-*, __N_0.distcp). Guarding those would deadlock rank 0 in all_gather and drop every non-zero rank's shard.

The DTensor v2 value worker reuses AutomodelCheckpointManager and is covered by the same change. The v1 path in native_checkpoint.save_checkpoint had the same bug and is fixed too.

Issues

Fixes a silent checkpoint hang on shared NFS storage. No linked issue.

Usage

No API or config change. Existing checkpointing calls are unaffected; only the number of processes writing the tokenizer directory changes (N to 1).

Before your PR is "Ready for review"

• Make sure you read and followed the Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? (see note below)
• Did you add or update any necessary documentation? (docstring explains the invariant and why it must not be applied to dcp.save)

Additional Information

Tests added:

• tests/unit/utils/test_native_checkpoint.py::TestSaveTokenizerOnRank0 covers: non-distributed writes; rank 0 writes and barriers; ranks 1/7/31 skip the write but still barrier; save_checkpoint() keeps dcp.save running on a non-zero rank while skipping the tokenizer.
• tests/unit/models/automodel/test_automodel_checkpoint.py::test_save_with_tokenizer_skipped_on_non_zero_rank pairs with the existing rank-0 test to cover the AutomodelCheckpointManager path.

Validation performed:

Check	Result
New tests (9 cases, incl. AutomodelCheckpointManager on ranks 0/5/31)	passed
Regression value: reverted the guard and re-ran	5 failed, confirming the tests actually catch this bug
ruff check and ruff format --check	clean

Note on local test runs: uv run currently fails to resolve dependencies in my environment for a reason unrelated to this change (nemo-rl[mcore] needs megatron-bridge, which pins transformers<=5.3.0, while nemo-rl requires >=5.5.0). The tests above were therefore executed against the real modules with a minimal stub for the optional transformers import. Please run the full suite in CI.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[dafu-wu]

/ok to test 27de939

[dafu-wu]

@yfw can you help review it?

[dafu-wu]

/ok to test 27de939

[yuki-97]

@jinglinglingling to review

[jinglinglingling]

LGTM. No concerns from my side.

[dafu-wu]

@jinglinglingling Thanks for your review! Could you also approve and run the four pending workflows?

[dafu-wu]

@jinglinglingling CI quality check is still stuck at Expected — Waiting for status to be reported. All currently triggered checks have completed, but I don't see a CI quality check run under the Checks tab. Could you help trigger it or check whether the required status check is stale?

[yfw]

/ok to test 38b62d3

[yfw]

/ok to test 8f38084