fix(parallel): shard SwiGLU weights in six model families

[chaofengw-nv]

Background

The Eagle VLM, GLM, Granite, OLMo, Qwen-VL, and StableLM tensor-parallel sharders reduce MLP width but leave SwiGLU gate/up/down weights unsharded. This passes full matrices to rank-local graph dimensions and causes reshape failures before inference. The defect is reproducible with small CPU arrays and does not require model downloads.

Exit Criteria

• Gate/up weights are column-sharded and down weights are row-sharded for TP 2/4/8; TP 1 remains unchanged.
• Rank shards reconstruct the original weights and their summed SwiGLU outputs match the unsharded calculation.
• Existing E2E acceptance criteria and GPU resource policy remain unchanged.

Implementation

Add the missing weight suffixes to each family's existing sharder, with one independently revertible commit per family. Each family owns its regression test; no cross-family model abstraction is introduced. There are no public API, ABI, dependency, or bundle-schema changes.

Change categories

• Model or runtime behavior
• Public API
• ABI
• Bundle or artifact format
• Dependencies
• Documentation only
• CI or developer tooling

Validation

Commands and Results

• python3 -m pytest families/eagle_vlm/tests/test_parallel_swiglu.py families/glm/tests/test_parallel_swiglu.py families/granite/tests/test_parallel_swiglu.py families/olmo/tests/test_parallel_swiglu.py families/qwen_vl/tests/test_parallel_swiglu.py families/stablelm/tests/test_parallel_swiglu.py -q --tb=short: 48 passed after the fix. The same regression cases produced 36 failures and 12 passes before the production change.
• python3 -m pytest families/eagle_vlm/tests families/glm/tests families/granite/tests families/olmo/tests families/qwen_vl/tests families/stablelm/tests -m 'not e2e and not gpu and not trt' -q --tb=short: 101 passed, 17 skipped; skipped tests require explicit E2E selection/checkpoints.
• python3 -m pytest tools/tests/test_architecture.py tools/tests/test_family_impact.py tools/tests/test_community_ci.py tools/tests/test_public_source_hygiene.py tools/tests/test_new_ci.py tools/tests/test_pr_metadata.py -q -p no:cacheprovider: 128 passed.
• ruff check families/{eagle_vlm,glm,granite,olmo,qwen_vl,stablelm}/parallel.py families/{eagle_vlm,glm,granite,olmo,qwen_vl,stablelm}/tests/test_parallel_swiglu.py: passed.
• python3 -m tools.model_ci validate: passed.
• git diff --check github/main...HEAD: passed.

Hardware, Environment, and Revisions

Tested tree: cbbe5c24324af75038e8090b036865696fa4221e, based on dddd2663336753c1102e47c281a925d96d899ac2. Linux x86_64, Python 3.12.3, NumPy 2.4.6, pytest 9.0.3; TensorRT 10.16.1.11 is available for imports. New tests use synthetic FP16/FP32 weights and CPU NumPy calculations, not GPU kernels or model checkpoints. CUDA execution was not exercised.

Not Run / Remaining Gaps

Target TensorRT 11.1 multi-GPU engine builds, inference, and checkpoint parity still require CI validation. CPU numerical reconstruction is not evidence of GPU E2E, performance, or model qualification. No premerge pass is claimed.

Contributor Self-Review

• I have completed a self-review of this change.

Checked family mapper weight layouts, rank-local graph dimensions, down-projection all-reduce placement, dtype/contiguity preservation, TP 1 behavior, and unchanged input arrays and replicated weights.

Notes For Future Readers

Review any family sharder together with its local regression test, then check the other family-local commits. Duplication here preserves independent family ownership and rollback. Rebuild affected TP engines to pick up the corrected weight partitioning; the serialized schema does not change. This fixes a deterministic shape defect, not every possible downstream multi-GPU failure.

Risk level

• Low
• Medium
• High

The production diff is small and CPU-reproducible, but it changes rank-local model weights in six families and still needs target-hardware validation.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2f1e07fc-d518-4deb-a953-f844ccbbee49

📥 Commits

Reviewing files that changed from the base of the PR and between cbbe5c2 and 0db4fb0.

📒 Files selected for processing (1)

• families/stablelm/parallel.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Summary

Summary

Fixes SwiGLU tensor-parallel sharding for Eagle VLM, GLM, Granite, OLMo, Qwen-VL, and StableLM.

• Shards w_gate and w_up by the last dimension.
• Shards w_down by the first dimension.
• Preserves TP 1 behavior.
• Supports TP sizes 2, 4, and 8.
• Adds family-owned regression tests for shapes, reconstruction, metadata, immutability, and numerical output equivalence.
• Covers float32 and float16.

Architecture impact

Status: PASS

• Family-owned files: Each model family has an independent parallel.py change and test_parallel_swiglu.py regression test.
• Shared surfaces: No shared API, ABI, dependency, bundle-schema, or E2E acceptance-criteria changes.
• Dependency direction: No new dependencies or new dependency directions.
• Affected consumers: Tensor-parallel SwiGLU weight loading and rank-local MLP inference for the six model families.
• Blast-radius questions: CI must confirm TensorRT 11.1 multi-GPU builds, inference, and checkpoint parity. Supplied validation reports passing regression, family, tooling, lint, and model checks.

Walkthrough

Changes

SwiGLU tensor-parallel sharding now covers w_gate, w_up, and w_down weights across Eagle VLM, GLM, Granite, OLMo, Qwen-VL, and StableLM. Parameterized tests validate partitioning and output reconstruction for multiple tensor-parallel sizes and dtypes.

SwiGLU sharding rules

Layer / File(s)	Summary
SwiGLU sharding classification families/eagle_vlm/parallel.py, families/glm/parallel.py, families/granite/parallel.py, families/olmo/parallel.py, families/qwen_vl/parallel.py, families/stablelm/parallel.py	w_down uses first-dimension sharding. w_gate and w_up use last-dimension sharding. Existing projection rules remain supported.
Parameterized sharding validation families/*/tests/test_parallel_swiglu.py	Tests cover tensor-parallel sizes 1, 2, 4, and 8 with float32 and float16. They validate shapes, dtypes, contiguity, metadata, reconstruction, immutability, and numerical equivalence.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

🚥 Pre-merge checks | ✅ 8 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 18 functions across 12 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (8 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Family Ownership Boundary	✅ Passed	PASS. The authoritative diff contains only two files in each of the six target family directories and no central registry or cross-family file. Each production change is local to its family sharder at…
Shared Semantic Neutrality	✅ Passed	PASS. The review-scoped diff contains only six family-owned parallel.py files and six family-owned regression tests under families/{eagle_vlm,glm,granite,olmo,qwen_vl,stablelm}. No shared code cha…
Benchmark Validation Integrity	✅ Passed	PASS. The production diff changes only tensor-parallel weight suffix classification in six sharders. It adds no benchmark, timing, metric, serialization, device-transfer, or synchronization measuremen…
Shared Change Blast Radius	✅ Passed	The pull request changes only six family-owned parallel.py modules and six matching family-owned regression test files. The authoritative diff contains no shared code, contract, tooling, example, be…
Title check	✅ Passed	The title clearly identifies the fix to SwiGLU weight sharding and names the six affected model families.
Description check	✅ Passed	The description covers the required background, exit criteria, implementation, change category, validation results, environment, remaining gaps, self-review, notes, and risk level. It does not link an…

---

_{Comment @coderabbitai help to get the list of available commands.}