Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
guozhihao-224 wants to merge 1 commit into from
Closed

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions docs/en/reference/tree_training.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,10 @@ actor:
| `pad_to_maximum` | bool | Yes | Must be `true` for tree training |
| `mb_spec.max_tokens_per_mb` | int | Yes | Max tokens per tree (must be set) |

NOTE: When tree training is enabled `max_tokens_per_mb` must be a multiple of
`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`

@gemini-code-assist gemini-code-assist Bot Apr 18, 2026

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`

divides 128, this simplifies to a multiple of 128.

## Implementation

Expand All @@ -72,7 +74,7 @@ Seq2: [A, G, H] [B] [G] tokens can attend to
The tree building process:

1. **Extract sequences**: Parse input_ids using attention_mask to get actual tokens
1. **Greedy packing**: Insert sequences into tries using first-fit decreasing strategy
1. **Greedy packing**: Insert sequences into tries using a first-fit strategy
1. **Trie compression**: Merge linear chains into single compressed nodes
1. **Mask generation**: Build block masks for efficient FlexAttention computation

Expand Down
Loading