This one is a real bug that never raised an error. LlamaBackbone.get_layout_map had the data-parallel and model-parallel positions swapped for the query kernel, the key/value kernels, and reverse_embeddings. Because LayoutMap tuples are just applied positionally, nothing complained; it silently produced the wrong sharding (e.g. the query kernel, shaped (hidden, num_query_heads, head_dim), got (model, data, None) instead of (data, model, None)). The fix restores the column/row-parallel convention that attention_output and the FFN kernels in the same function already use, and leaves GQA k/v replicated.
Comes with a Llama mesh-shape sweep (7B non-GQA, 70B 8:1 GQA) and a live-preset sweep. Part of #2807, fixed in #2815.
This one is a real bug that never raised an error.
LlamaBackbone.get_layout_maphad the data-parallel and model-parallel positions swapped for the query kernel, the key/value kernels, andreverse_embeddings. BecauseLayoutMaptuples are just applied positionally, nothing complained; it silently produced the wrong sharding (e.g. the query kernel, shaped(hidden, num_query_heads, head_dim), got(model, data, None)instead of(data, model, None)). The fix restores the column/row-parallel convention thatattention_outputand the FFN kernels in the same function already use, and leaves GQA k/v replicated.Comes with a Llama mesh-shape sweep (7B non-GQA, 70B 8:1 GQA) and a live-preset sweep. Part of #2807, fixed in #2815.