Skip to content

perf(sft): use modified first-fit decreasing packing - #2369

Open
dyurk-lila wants to merge 4 commits into
mainfrom
mffd-sequence-packing
Open

dyurk-lila wants to merge 4 commits into
mainfrom
mffd-sequence-packing

Conversation

@dyurk-lila

@dyurk-lila dyurk-lila commented Oct 1, 2026 •

Copy link
Copy Markdown
Collaborator

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.

PackedDataCollator selects 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:

  • Chat: packing scores tied in all 800 comparisons, though assignments differed.
  • Math proofs: MFFD won 115 comparisons, FFD won 77, and 608 tied. MFFD removed one bin in five comparisons and never used more bins.

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:

pytest -q \
  tests/backends/skyrl_train/distributed/test_bin_packing.py \
  tests/train/test_sft_packing_collate.py \
  tests/train/test_collation_vectorization_equivalence.py \
  tests/train/test_sft_config.py

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.

  1. perf(sft): use modified first-fit decreasing packing #2369 — use modified first-fit decreasing packing
  2. perf(train): make sequence packing alignment-aware #2370 — make sequence packing alignment-aware
  3. refactor(megatron): share packed segment layout derivation #2371 — share packed segment layout derivation
  4. perf(sft): align packed bin counts to data parallelism #2372 — align packed bin counts to data parallelism
  5. perf(megatron): autotune fused LM-head Triton kernels #2373 — autotune fused LM-head Triton kernels
  6. feat(megatron): fuse LM-head entropy for RL #2374 — fuse LM-head entropy for RL
  7. feat(megatron): add block-sparse fused LM-head forward #2375 — add block-sparse fused LM-head forward
  8. perf(megatron): add block-sparse fused LM-head backward #2376 — add block-sparse fused LM-head backward
  9. perf(megatron): fuse sample-support candidate projection with Triton #2521 — fuse sample-support candidate projection with Triton

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 selectable PackingStrategy.

Adds a ModifiedFirstFitDecreasing packer (size-class phases, bisect-based fitting, first-fit leftovers) and wires PackedDataCollator to PackingStrategy.MODIFIED_FIRST_FIT_DECREASING. Docs and SFTConfig comments 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.

@dyurk-lila
dyurk-lila added this pull request to stack #2377 October 1, 2026 14:10

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@greptile-apps

greptile-apps Bot commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 4/5

[Medium risk] Switches sequence packing algorithm from FFD to MFFD.

The PR appears safe to merge, with a non-blocking concern about controller packing time for large batches.

Findings

  1. P2 Repeated list removals slow packing ▶
Diagram
%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Sequence lengths] --> B[Classify and sort]
  B --> C[Place medium and small items with large items]
  C --> D[Fill existing bins from leftovers]
  D --> E[First-fit remaining items into new bins]
  E --> F[Adjust bin count for DP]
  F --> G[Build and shard packed rows]
Loading

Reviews (1) · Last reviewed commit: "Address MFFD review feedback"

Comment on lines +303 to +304
item = remaining_items.pop(chosen_index)
negative_remaining_sizes.pop(chosen_index)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 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!

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>
@dyurk-lila
dyurk-lila force-pushed the mffd-sequence-packing branch from ee7b212 to cf5d4e6 Compare October 23, 2026 16:17
@avigyabb
avigyabb removed this pull request from stack #2377 October 23, 2026 22:57

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants