Skip to content

fix(integrations): forward distributed sharding options - #41

Merged
jfischoff merged 2 commits into
mainfrom
feat/expose-distributed-sharded
Aug 17, 2026
Merged

jfischoff merged 2 commits into
mainfrom
feat/expose-distributed-sharded

Conversation

@jfischoff

Copy link
Copy Markdown
Contributor

Summary

FlashPack's core loader already supports both rank-0 read+broadcast and one-shard-per-rank read+replicate for arbitrary packed dtypes. The high-level integrations did not expose that contract consistently:

  • Diffusers pipelines silently dropped distributed_sharded and distributed_shard_strategy, so callers got rank-0 broadcast even when they requested sharded reads.
  • Transformers and Diffusers model loaders happened to pass the options through **kwargs, but did not declare or explicitly forward them.

This change makes both options first-class across all three high-level paths and documents that distributed loading works for ordinary BF16/FP32 packs as well as quantized packs.

Changes

  • Forward sharded-load controls through FlashPackDiffusionPipeline and load_sub_model_flashpack.
  • Declare and explicitly forward the controls in both model integrations.
  • Add dependency-free AST guardrails that fail if any high-level layer drops either option again.
  • Add a high-level distributed loading example to the README.

Test plan

  • uv run --extra dev pytest -q -m 'not network' — 244 passed, 4 skipped, 4 deselected.
  • .venv/bin/pre-commit run --all-files — all hooks passed.
  • Full default test selection also passed 246 tests; its only two failures were pre-existing network integration tests importing optional diffusers/transformers dependencies that are not in the dev extra.

Expose the core distributed_sharded and shard strategy controls through
Transformers, Diffusers models, and Diffusers pipelines.

Add dependency-free API plumbing guards and document that the distributed
loader applies to ordinary BF16 and FP32 packs as well as quantized packs.
@jfischoff

Copy link
Copy Markdown
Contributor Author

@codex review

@jfischoff
jfischoff merged commit e444064 into main Aug 17, 2026
6 checks passed
@jfischoff
jfischoff deleted the feat/expose-distributed-sharded branch August 17, 2026 17:05
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