torchtitan construiu o HSDP + suporte diloco, provavelmente é o melhor lugar agora para começar a fazer pesquisa em aprendizado descentralizado. Também vem com suporte para muitas arquiteturas (llama3, llama4, deepseekv3...) assim como todo o possível paralelismo (6d?). A equipe do Pytorch cozinhou aqui.