[Prototype] Concatenated weights and linear layers

[jlamypoirier]

✨ Description

Handle concatenated weights, i.e. weights that are stored as one in memory for computation / optimization purpose but actually consist of multiple logically distinct ones. Ex. Key and value, gated MLP layer 1, MoE layers, SSM in_proj, etc.

This lets the engine know about the structure, and allow configuring each sub-parameter independently (lr_scale, initialization, weight decay, peft enabling. Bias enabling split possible but not implemented).

Prototype is working but postponing finalization in favor of higher priority content. Remaining:

• Simplify linear layer concatenation and make it safer (especially peft)
• Use for remaining linear layers
• Make remaining tests pass.