fix: shard single-row gate weights along columns for Muon - #3723

fix: shard single-row gate weights along columns for Muon#3723

hubert-marek merged 3 commits into

mainPrimeIntellect-ai/prime-rl:mainfrom

fix/muon-single-row-fsdp-shardingPrimeIntellect-ai/prime-rl:fix/muon-single-row-fsdp-shardingCopy head branch name to clipboard