Um contribuidor adicionou suporte nativo ao `nemotron_h_puzzle` no mlx-lm por meio do PR #1535, permitindo que o modelo Nemotron Puzzle 75B seja executado em um M2 Max com 64GB. O trabalho envolveu corrigir o suporte à configuração em blocos, o remapeamento de tensores e alinhar o comportamento da normalização/roteador FP32 da NVIDIA.

  • A quantização de especialistas em 4 bits (com camadas densas de 6 bits e cabeça de saída BF16) utiliza um tamanho de checkpoint de 42,03 GiB e atinge um pico de 49,68 GB de memória.
  • A configuração de 4 bits alcança uma velocidade média de geração de 14,27 tok/s, comparado a 10,53 tok/s para a variante de 5 bits.
  • As verificações de tarefas locais pontuaram 24/30 para 4 bits contra 21/30 para 5 bits, enquanto a recuperação de contexto longo foi 4/5 contra 0/5.
  • Um bug que causava discrepâncias na saída SSM da primeira camada foi resolvido ao alinhar o limite de promoção do BF16 para FP32 da NVIDIA, melhorando a similaridade cosseno para 0,999998.
  • A cabeça de saída BF16 é obrigatória; quantizar a lm_head com vocabulário de 131k em 4 bits produziu lixo repetitivo.

O autor observa que o checkpoint de 5 bits roda próximo ao limite de memória de 64GB, sugerindo que configurações com maior memória podem produzir resultados diferentes.