Repository navigation
perf(sft): use modified first-fit decreasing packing - #2369
dyurk-lila wants to merge 4 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces the Modified First-Fit Decreasing (MFFD) bin-packing algorithm as the default sequence packing strategy, replacing the previous First-Fit Decreasing (FFD) implementation. The changes include adding the ModifiedFirstFitDecreasing class, updating the configuration, documentation, and collators to use MFFD, and expanding the test suite to verify the correctness of the new packing strategy. I have no feedback to provide as there are no review comments.
|
| item = remaining_items.pop(chosen_index) | ||
| negative_remaining_sizes.pop(chosen_index) |
There was a problem hiding this comment.
Repeated list removals slow packing
For large global SFT batches, this loop repeatedly removes items from the middle of remaining_items. Each removal shifts the rest of the Python list, so packing can spend substantial time moving items on the controller every training step. Consider avoiding repeated middle-of-list removals and benchmarking packing time at the 2,048-example batch size used in the PR’s packing comparisons.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
fe02ff9 to
9dcca01
Compare
97f9aad to
ee7b212
Compare
Signed-off-by: Dominic Yurk <dyurk@lila.ai>
Signed-off-by: Dominic Yurk <dyurk@lila.ai>
Signed-off-by: Dominic Yurk <dyurk@lila.ai>
Signed-off-by: Dominic Yurk <dyurk@lila.ai>
ee7b212 to
cf5d4e6
Compare
What does this PR do?
Use Modified First-Fit Decreasing (MFFD) for controller-level SFT sequence packing, retaining FFD as a generic packing strategy.
How it works
MFFD classifies sequences by size, places large sequences and compatible smaller sequences in phases, then assigns remaining sequences with true first-fit placement in bin creation order. Binary search selects fitting items from the sorted leftovers; the final placement pass uses a direct Python scan.
PackedDataCollatorselects MFFD when sequence packing is enabled. Shared packer tests run against both FFD and MFFD, and collation tests verify MFFD’s expected row assignments.Packing results
Recorded benchmarks used Nemotron-3.5-Lightning SFT data, global batch size 2048, and 32K bin capacity. Across 800 alignment/scenario comparisons per slice:
Across all 1,600 comparisons, those five reductions average 0.003125 fewer 32K bins per batch. These are modest packing improvements, not an end-to-end training speedup claim.
The implementation follows the original MFFD paper and is informed by NeMo RL’s packing algorithms.
Validation
The following CPU test selection passed 236 tests, with one CUDA-only test skipped:
Ruff, Black, and secret-detection checks passed. GPU tests were not rerun after rebasing onto current
main.Based on
main; #2370 adds alignment-aware capacity accounting above this algorithm layer.Review order
Part 1 of 9 in a linked series. The stack is based on
main; each later PR targets the preceding layer. Review and merge from bottom to top.Review only this layer.
Note
Medium Risk
Changes which sequences share each packed micro-batch row and can alter loss grouping across bin boundaries, though DP bin-count padding and collator invariants stay the same; FFD remains available for explicit use.
Overview
Switches Megatron SFT controller-level sequence packing from First-Fit Decreasing (FFD) to Modified First-Fit Decreasing (MFFD) when
use_sequence_packing=true, while keeping FFD as a selectablePackingStrategy.Adds a
ModifiedFirstFitDecreasingpacker (size-class phases, bisect-based fitting, first-fit leftovers) and wiresPackedDataCollatortoPackingStrategy.MODIFIED_FIRST_FIT_DECREASING. Docs andSFTConfigcomments now describe MFFD bin capacity; shared packer tests cover both algorithms plus MFFD-specific assignment cases.Reviewed by Cursor Bugbot for commit 97f9aad. Bugbot is set up for automated code reviews on this repo. Configure here.