Skip to content

Make the encode pool opt-out behind a parallel feature - #2289

Open
ArthurZucker wants to merge 1 commit into
feat/multi_threaded_pipelinefrom
feat/rayon-optional
Open

Make the encode pool opt-out behind a parallel feature#2289
ArthurZucker wants to merge 1 commit into
feat/multi_threaded_pipelinefrom
feat/rayon-optional

Conversation

@ArthurZucker

Copy link
Copy Markdown
Collaborator

parallel joins the default feature set, so nothing changes unless you ask for it. With --no-default-features the private rayon ThreadPool is never built, no worker threads are spawned, and the job-splitting machinery is not compiled — encode takes the serial path a small batch already took.

66,272 bytes on the stripped binsize_pipeline example: 2,124,096 → 2,057,824. Measured with progressbar held on, so the number is the pool rather than indicatif.

What this is not

It does not make rayon optional, and the title deliberately doesn't claim it does. rayon stays a hard dependency with the feature off, because ptr_hash declares it unconditionally to build the MPHF — and rayon-cond and rdst pull it as well:

rayon v1.12.0
├── ptr_hash v2.0.2   -> tk-encode
├── rayon-cond v0.4.0 -> tk-encode
└── rdst v0.20.14

cargo tree still shows rayon either way. What the flag buys is a build with no thread pool and no parallel encode path, which is what matters for embedded and single-threaded hosts. It is not a dependency reduction, and a reviewer should not expect one.

Making rayon genuinely optional would mean serial fallbacks at all 19 maybe_par_* call sites (7 in tk-encode, 12 in tk-train) plus tk-train opting in — and it still wouldn't remove rayon, for the reason above. That work also lands squarely in tokenizer/mod.rs, encoding.rs and padding.rs, which the pipeline-only refactor replaces. Not worth doing unless ptr_hash changes, and if it does it should be re-derived from scratch.

Implementation notes

  • The parallel batch body moves out of encode into a gated encode_parallel. It takes the owned input storage rather than a slice because JobCore ends up owning it, which is also why refs is rebuilt inside.
  • With the feature off, the striding and segment-planning helpers have no caller. They're left compiled-but-unused under a scoped allow(dead_code) rather than cfg'd one by one: they reference each other densely and LTO drops all of it from the binary regardless, so cfg-ing each would be churn in files this branch's successor replaces.
  • Nested if let rather than a let-chain in the chunk-commit scatter — this crate is edition 2018.

Verification

default --no-default-features
build clean, 0 warnings clean, 0 warnings
lib tests 291 passed 290 passed¹
parallel oracle 15 passed 15 passed
doc tests 20 passed 20 passed

¹ the pool's own test is gated with the module.

…ture

`parallel` is added to the default feature set, so nothing changes unless you ask.
With `--no-default-features` the private rayon `ThreadPool` is not built, no worker
threads are spawned, and the job-splitting machinery is not compiled: `encode` takes
the serial path that a small batch already took.

Worth 66,272 bytes on the stripped `binsize_pipeline` example, 2,124,096 ->
2,057,824, measured with `progressbar` held on so the number is the pool and not
indicatif.

**This does not make rayon optional, and the title deliberately does not say it
does.** rayon remains a hard dependency with the feature off, because `ptr_hash`
declares it unconditionally to build the MPHF, and `rayon-cond` and `rdst` pull it
too:

    rayon v1.12.0
    |-- ptr_hash v2.0.2  -> tk-encode
    |-- rayon-cond v0.4.0 -> tk-encode
    |-- rdst v0.20.14

So `cargo tree` still shows rayon either way. What the flag buys is a build with no
thread pool and no parallel encode path, which matters for embedded and
single-threaded hosts; it is not a dependency reduction.

Making rayon genuinely optional would need serial fallbacks at all 19
`maybe_par_*` call sites (7 here, 12 in tk-train) plus `tk-train` opting in — and it
would still not remove rayon, for the reason above. That refactor also lands in
`tokenizer/mod.rs`, `encoding.rs` and `padding.rs`, which the pipeline-only work
replaces. Not worth doing unless `ptr_hash` changes, at which point it should be
re-derived rather than resurrected.

Two implementation notes. The parallel batch body moves out of `encode` into a gated
`encode_parallel`, which takes the owned input storage rather than a slice because
`JobCore` ends up owning it. And with the feature off, the striding and segment
planning helpers have no caller; they are left compiled-but-unused under a scoped
`allow(dead_code)` rather than cfg'd one by one, because they reference each other
densely and LTO drops all of it from the binary regardless — cfg-ing each would be
churn in files the pipeline-only refactor is replacing.

Tests pass both ways: 291 lib tests with the feature, 290 without (the pool's own
test is gated with the module), plus 15 parallel-oracle and 20 doc tests unchanged.
Both configurations build warning-clean.
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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.

2 participants