Skip to content

feat: add dense MLP activation sharding#550

Open
gauravsaini wants to merge 1 commit into
JustVugg:devfrom
gauravsaini:dense-sharding-upstream
Open

feat: add dense MLP activation sharding#550
gauravsaini wants to merge 1 commit into
JustVugg:devfrom
gauravsaini:dense-sharding-upstream

Conversation

@gauravsaini

Copy link
Copy Markdown

Summary

  • add a coordinator-to-worker TCP activation protocol for contiguous dense MLP layers
  • add coli cluster dense-worker and --dense-shards host:port:first:last
  • keep attention, KV cache, routing, and token generation on the coordinator
  • load only the selected dense gate/up/down tensors on each worker
  • add wire-level protocol/ownership tests and a token-exact CPU teacher-forcing gate

Correctness gate

The optional c/tests/test_dense_sharding.py test starts a worker and compares the exact TF=1 oracle mismatch signature between local CPU execution and delegated dense execution. It skips cleanly when the gitignored glm_tiny fixture is unavailable.

Native verification on macOS:

  • make colibri
  • make test-c
  • ./tests/test_dense_protocol
  • dense sharding correctness test: skipped only because glm_tiny is not present locally

This is intentionally limited to dense MLP activation sharding. Expert-worker and WebGPU work remain separate PRs. The existing bundled PR #380 remains open while the work is split.

@JustVugg

Copy link
Copy Markdown
Owner

This is the strongest of your three (it ships test_dense_sharding.py, a real token-exact baseline-vs-delegated gate — exactly the discipline the others need). But same honest note: dense MLP sharding across nodes is breadth, and the priority right now is depth on single-box models, so I'm deferring rather than merging today. It's env-gated OFF (default path safe) and it conflicts with #551/#552 anyway (all three edit the same colibri.c lines). Kept open — if any of the trio lands when breadth is the goal, this is the one with the correctness gate, so it'd be first. Thank you for the clean split and the test.

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