Fix FPDT loss layout for batch-first masks - #8541
taking-lying-flat wants to merge 1 commit into
Conversation
FPDT returns gathered losses as [sequence, batch], while its loss mask uses [batch, sequence]. Flattening both silently weights the wrong tokens when a batch contains multiple examples and a nonuniform mask. Return contiguous batch-first losses. In backward, select the local sequence slice and transpose it back for the existing chunked backward. Validated locally against dense cross entropy on CPU/gloo at SP=1/2 and on an RTX A1000 Laptop GPU at SP=1, comparing loss, parameter gradients, and one SGD update. Regression tests remain outside this production-only change. Changed-file pre-commit checks passed. Signed-off-by: taking-lying-flat <1615405@qq.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![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: 74febd0228
ℹ️ 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".
|
|
||
| return loss_all | ||
| # Match the batch-first loss mask produced by FPDT_InputConstruct. | ||
| return loss_all.t().contiguous() |
There was a problem hiding this comment.
This is a non-merge commit, but its metadata contains no Signed-off-by trailer; the repository requires every non-merge commit to be created with --signoff, so this commit does not satisfy the contribution requirements and should be recreated with the configured Codex <codex@openai.com> identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
1 participant
FPDT silently applies loss masks to the wrong tokens when the micro-batch contains multiple examples and the mask is nonuniform.
FPDT_LogitsLossreturns gathered losses as[sequence, batch], whileFPDT_InputConstructreturns its mask as[batch, sequence]. The Megatron-DeepSpeed FPDT caller passes these losses directly to a flattened masked reduction, so the mismatch does not raise a shape error.Return contiguous
[batch, sequence]losses. In backward, select the rank's slice along the sequence dimension and transpose it back before the existing chunked gradient computation. The mask remains global to match the gathered losses.This PR changes two executable lines and adds one explanatory comment in
fpdt_layer.py.Validation was run locally; test changes are not included in this PR:
Multi-GPU NCCL was not tested on this single-GPU machine.