fix(init): initialize DTensor-sharded models whose embedding has a padding_idx

[yisongbetter]

What does this PR do ?

Fixes Checkpointer.initialize_model_weights skipping the model's random init when an
nn.Embedding with padding_idx is already a DTensor. For custom models materialized after
sharding (from_config + FSDP2 / PP) every parameter then stayed as the uninitialized memory
to_empty() handed out.

Symptom on Kimi-K3 (examples/llm_benchmark/kimi/kimi_k3_gb200.yaml): the 2-node benchmark trains an
all-zero model (loss pinned at ln(163840) = 12.0066, grad_norm 0.0000, every RMSNorm weight 0), the
64-node benchmark a garbage one (grad_norm nan at step 0, loss nan from step 1). Both were verified
by printing per-parameter statistics right after setup(): 636/636 parameters all-zero on every rank.

Root cause: #1675 added has_padding_idx (embedding with padding_idx whose weight is a DTensor) to
skip_initialize_weights to avoid HF's DTensor-unsafe weight[padding_idx] = 0. The comment describes
clearing padding_idx around the init and zeroing the row locally; the code skipped the whole init
instead. Models that build their embedding with padding_idx and rely on the Checkpointer's random init are
affected (kimi_k3, kimi_linear, baichuan; ernie4_5, glm5_next, inkling, laguna opt out of
that init via _skip_init_weights_on_load); models that do not pass it (DeepSeek V3/V3.2/V4, GLM4-MoE,
GPT-OSS, …) never hit the skip, which is why the other llm_benchmark configs were fine.

With the padding row now correctly zero, the mock benchmark data should not contain the padding id
either: a randomly placed padding token is a zero vector at an arbitrary position (including sequence
starts), a shape real data never has, and on very deep models it can make long random-data runs
numerically noisy. MockIterableDataset therefore gains an exclude_token_ids option and the Kimi-K3
benchmark config excludes id 0.

Changelog

• nemo_automodel/shared/embedding_padding.py (new): zero_embedding_row_(weight, row) zeroes one
  embedding row on the rank-local shard for plain tensors and Replicate/Shard DTensors, without
  integer-indexing the DTensor (compute_local_shape_and_global_offset, same approach as
  training/embedding_row_repair.py).
• nemo_automodel/components/checkpoint/checkpointing.py: drop has_padding_idx from the skip list;
  clear padding_idx on DTensor embeddings for the duration of initialize_weights() (so HF's
  _init_weights skips its indexing op), restore it in a finally, then zero the padding row locally.
• nemo_automodel/components/models/kimi_k3/model.py, kimi_linear/model.py: init_weights zeroes the
  padding row through the helper instead of weight[padding_idx].zero_() on a possibly sharded weight.
• nemo_automodel/components/datasets/llm/mock_iterable_dataset.py: exclude_token_ids on
  MockIterableDataset / MockIterableDatasetConfig — uniform sampling over the vocabulary minus the
  listed ids; default behaviour unchanged. examples/llm_benchmark/kimi/kimi_k3_gb200.yaml sets
  exclude_token_ids: [0] (KimiK3TextConfig.pad_token_id).
• examples/llm_benchmark/kimi/kimi_k3_gb200.yaml header: figures measured with this fix applied
  (comments only; the earlier figures came from the uninitialized model).
• Tests: tests/unit_tests/shared/test_embedding_padding.py,
  tests/unit_tests/models/kimi_k3/test_init_weights_dtensor.py,
  TestInitializeModelWeightsPaddingIdxDTensor in tests/unit_tests/checkpoint/test_checkpointing.py
  (single-rank gloo CPU DTensors; asserts init runs, padding_idx is cleared during init and restored
  even when init raises, padding row zero, other rows initialized);
  TestMockIterableDatasetExcludeTokenIds in tests/unit_tests/datasets/llm/test_mock_iterable_dataset.py.

Validation

• Unit tests above + the existing TestInitializeModelWeights cases: 26 passed inside the nemo-automodel:26.06 container (single-rank gloo, CPU DTensors).
• Weight statistics after setup() for the 2-node Kimi-K3 mini benchmark (from_config, FSDP2):
  before this PR: 636/636 parameters all-zero on every rank (RMSNorm weights included). After: 18 all-zero parameters per rank — exactly the dt_bias of the 18 KDA layers, which init_weights zeroes by design; embed_tokens std 2.0e-2 (initializer_range), A_log mean 1.97 (uniform(1, 16).log()), RMSNorm weights 1, no NaN/Inf.
• 12-step 2-node Kimi-K3 mini benchmark (24 layers / 64 experts, hybridep, FSDP2) on the fixed tree: loss 12.489 → 12.466 over 12 steps, finite throughout (before: pinned at 12.0066 with zero gradients); recipe Average MFU 20.2%.
• tests/unit_tests/datasets/llm/test_mock_iterable_dataset.py: 23 passed in the container (18 existing + 5 new).
• 2-node Kimi-K3 (93 layers / 4 experts, pp4×ep2, mock data): with the padding id kept out of the data,
  6 steps finite (12.4945 → 12.4850); the same run with every sequence starting on the padding id lost
  numerical stability after the first optimizer step — the case the new option removes.
• ruff check + ruff format --check clean (0.12.12).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (three test files, listed above)
• Did you add or update any necessary documentation? (docstrings on the helper and the init path; no docs page covers weight initialization)

Additional Information

• Related to fix: skip embedding[padding_idx] = 0 with TP #1675 (introduced the skip) and feat: add Kimi K3 model support #3259 (Kimi-K3, whose embedding carries padding_idx).
• The llm_benchmark Kimi-K3 numbers published so far were measured on this uninitialized model; they
  are being re-measured on top of this fix (perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 / perf(distributed): ring-pooled PP recv buffers and Kimi-K3 2k benchmark shape #3780 will carry the new numbers).

[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.

[yisongbetter]

/ok to test ea125d4

[yisongbetter]

/ok to test ad8e8f1

[yisongbetter]

/ok to test e4000ab