Skip to content

docs: fix tree packing algorithm name and token alignment constraint#1205

Closed
guozhihao-224 wants to merge 1 commit into
areal-project:mainfrom
guozhihao-224:docs/fix-tree-training-docs
Closed

docs: fix tree packing algorithm name and token alignment constraint#1205
guozhihao-224 wants to merge 1 commit into
areal-project:mainfrom
guozhihao-224:docs/fix-tree-training-docs

Conversation

@guozhihao-224

Copy link
Copy Markdown
Collaborator

Summary

Two corrections to the tree training reference documentation:

  1. Packing algorithm: Change "first-fit decreasing" to "first-fit". The code in _greedy_build_tries() iterates sequences in original order without sorting by decreasing size. (A proper FFD implementation exists in areal/utils/seqpack.py for trajectory redistribution, but tree packing intentionally uses plain first-fit to preserve prefix locality.)

  2. Token alignment constraint: Document the full constraint lcm(BLOCK_SIZE, parallel_size) instead of just BLOCK_SIZE (128). In most configurations where parallel_size divides 128 this equals 128, but the actual code at tree.py:341 uses math.lcm() which can be stricter with unusual parallel configs.

Verification

  • Confirmed no sort step in _greedy_build_tries() at tree.py:476-523
  • Confirmed math.lcm(BLOCK_SIZE, parallel_size) at tree.py:341
  • Git history shows no sort step was ever present or removed

Test plan

  • pre-commit run --all-files passes
  • ./docs/build_all.sh builds without errors

Ref: #1165 (Findings #6, #7)

🤖 Generated with Claude Code

1. Change "first-fit decreasing" to "first-fit". The code in
   _greedy_build_tries() iterates sequences in original order without
   sorting by decreasing size.

2. Document the full token alignment constraint:
   lcm(BLOCK_SIZE, parallel_size) instead of just BLOCK_SIZE (128).
   When parallel_size divides 128 (common case) this equals 128,
   but can be larger with unusual parallel configs.

Ref: areal-project#1165 (Findings areal-project#6, areal-project#7)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@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 updates the tree training documentation to clarify that max_tokens_per_mb must be a multiple of the least common multiple of the block size and parallel size. It also updates the description of the greedy packing strategy from first-fit decreasing to first-fit. Feedback was provided to generalize the definition of parallel_size to ensure accuracy across different backend engines like Megatron and Archon.

`BLOCK_SIZE` (128).
NOTE: When tree training is enabled, `max_tokens_per_mb` must be a multiple of
`lcm(BLOCK_SIZE, parallel_size)`, where `BLOCK_SIZE` is 128 and
`parallel_size = tp_size × sp_size`. In most configurations where `parallel_size`

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.

medium

The definition of parallel_size as tp_size × sp_size is accurate for the FSDP backend, but the Megatron and Archon engines currently only pass the tensor parallel degree (tp_size) to the packing function. To ensure the documentation is accurate across all supported backends, consider clarifying that parallel_size represents the product of parallelism dimensions that require block alignment for the specific engine being used.

Suggested change
`parallel_size = tp_size × sp_size`. In most configurations where `parallel_size`
`parallel_size` (the product of parallelism dimensions requiring block alignment, e.g., `tp_size × sp_size` for FSDP). In most configurations where `parallel_size`

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant