Repository navigation
Conversation
… metrics Signed-off-by: khazic <khazzz1c@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c4dae9188c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| supervised = outputs.eval_mask.to(torch.float32) * outputs.block_keep_mask.to(torch.float32).unsqueeze(-1) | ||
| weight = _loss_weight_mask(supervised, block_size, loss_decay_gamma) # [B, N, L] |
There was a problem hiding this comment.
When a dropped block contains padded out-of-range target IDs or non-finite draft/teacher logits, constructing a zero supervised weight does not actually exclude it: cross-entropy and softmax/BCE still evaluate every position before weighting, so CE can raise on the discarded label and NaN * 0 can poison the loss, gradients, and metrics. Select or replace unsupervised inputs with safe values before these nonlinear operations rather than relying only on post-hoc multiplication.
Useful? React with 👍 / 👎.
| block_keep_mask: Kept-anchor mask ``[batch, num_blocks]``; dropped anchors | ||
| are excluded from the loss and the acceptance metrics. |
There was a problem hiding this comment.
The reviewed commit has no Signed-off-by trailer, so it violates the repository's mandatory signed-commit policy and can be rejected by the DCO gate; recreate the commit using git commit -s before merging.
AGENTS.md reference: AGENTS.md:L39-L43
Useful? React with 👍 / 👎.
|
Superseded by the upstream PRs (NVIDIA-NeMo/Megatron-Bridge NVIDIA-NeMo#5241 NVIDIA-NeMo#5242 NVIDIA-NeMo#5243 NVIDIA-NeMo#5244). |
1 participant
What
block_keep_maskwas consumed only by thetaudiagnostic. The loss weight was built fromeval_maskalone, so a dropped anchor block (keep = 0) whose eval prefix was still set received full CE/TV/confidence gradient, and its positions were averaged intoaccept_rate, biasing both the objective and the diagnostics.Fix
Define the supervised-position mask once in
dspark_loss(supervised = eval_mask * block_keep_mask.unsqueeze(-1)), derive the decayed loss weight from it, and pass it into_acceptance_metrics(now keyword-only per the repo style rule on same-type parameters). One mask drives the loss weight,accept_rate, and thetaublock validity, so the three sites cannot drift.Tests
Two new unit tests: dropped blocks produce exactly zero gradient and the loss equals the eval-mask-zeroed equivalent;
accept_rate/tauignore a dropped block that deliberately mismatches the teacher. Both fail on the parent commit and pass here. All 18 dspark unit tests pass; ruff is clean.Merge note
khazic/fix/dspark-masked-input-robustnesssanitizes padded inputs keyed oneval_maskalone. Whichever of the two branches merges intokhazic/feat/dspark-coresecond should rebase that sanitization onto the combinedsupervisedmask (one-line change: buildsupervisedfirst and usesupervised.bool()as the scrub mask), so that positions of dropped anchors are scrubbed as well as zero-weighted.