WeightlessLinear

A drop-in replacement for max.nn.Linear that reconstructs its weight in-kernel, and the stacked variant that keeps GQA projections distinct.

core/weightless_linear.py provides the layer that replaces max.nn.Linear throughout a weightless model.

WeightlessLinear

WeightlessLinear accepts the same constructor signature as max.nn.Linear (in_dim, out_dim, dtype, device, has_bias, quant_config, name, …) plus a lora_rank, and absorbs any extra keyword arguments. That makes functools.partial(WeightlessLinear, ...) a drop-in for the linear_cls parameter that every standard MAX architecture threads through construction.

Instead of x @ weight.T, its __call__ computes:

y = s · (W0 x) + U (Vᵀ x)

via either the generated_linear custom op or the GEMM formulation, auto-selected by accelerator and lora_rank (GEMM on NVIDIA GPUs at any rank; on Metal, GEMM when lora_rank >= 8). The full [out_dim, in_dim] weight is never materialized.

Owned vs injected weights

The per-layer theta weights are Weight objects owned by the layer, so their state-dict names are <name>.s, <name>.U, <name>.V:

The layer doesn’t own the shared block W0. The model declares it once (one per family) and injects it before the forward pass:

layer.set_w0(w0)   # must be called before __call__

Calling the layer without an injected W0 raises RuntimeError.

Compatibility surface

To satisfy framework code that inspects a linear layer, WeightlessLinear exposes a placeholder .weight with the correct [out_dim, in_dim] shape (its value is never used for compute), plus .bias, .weight_scale, .input_scale, and .weight_scale_2 (all None on the unquantized weightless path). It implements the Shardable protocol so sharding_strategy reads and writes don’t crash, but .shard() raises NotImplementedError, because tensor parallelism doesn’t yet compose with in-kernel reconstruction.

WeightlessStackedLinear

The base StackedLinear concatenates its children’s .weight tensors into one stacked weight and calls linear once. That path materializes the full weight and never invokes the custom op, so it’s incompatible with weightless children.

WeightlessStackedLinear overrides __call__ to iterate the children and call each one directly, so every projection runs its own reconstruction and returns a list of per-child outputs. This keeps GQA configurations correct: when q_proj, k_proj, and v_proj have different output dimensions, each output keeps its own out_dim rather than being obscured by concatenation. Child weights appear at self_attn.q_proj.* (not qkv_proj.q_proj.*), matching the checkpoint layout without a per-architecture weight mapping.