Pool M3 training sparse weights without token-vocabulary expansion

[Afloat16]

M3 sparse training currently materializes a batch × sequence × vocabulary tensor before reducing repeated token occurrences. For long sequences and a large vocabulary, this allocation dominates the sparse branch and can prevent otherwise feasible fine-tuning batches.

Pool directly into the batch × vocabulary result. Select the first maximum occurrence using detached index reductions, then gather from the original differentiable token weights. This preserves the existing first-maximum gradient convention, including positive ties, absent/special tokens and NaN propagation, while requiring storage proportional to batch × (sequence + vocabulary).

The 11 actual CPU regression cases use a genuine local BertModel and an independent token-dictionary oracle. Before this change, the two allocation checks fail and the other nine pass; after this change all 11 pass. The original import/trainer compatibility modules also pass all nine unchanged tests in both runs. Default discovery collects the same 32 items without errors; the remaining pretrained-model integration tests were collected and were not executed. No model or dataset download is needed for the executed tests.

Validated with Python 3.12.10, Torch 2.14.1+cpu and Transformers 5.18.0. Profiler bounds cover small reproducible CPU allocations; production-scale GPU memory, throughput and training quality were not benchmarked. Existing evaluation aggregation, local-negative objectives and passage MaxSim masking are unchanged.