Restauration de ruGPT-3 XL pour transformers

Conversion et restauration de ruGPT-3 XL pour les frameworks modernes

Les développeurs ont longtemps eu du mal à faire fonctionner le modèle linguistique ancien ai-forever/rugpt3xl (1,3 milliard de paramètres) sorti en 2021. Initialement entraîné depuis zéro sur un corpus russe comme une variante GPT-2, il est stocké au format checkpoint Megatron-LM, nécessitant PyTorch 1.7 et transformers 3.5. Le convertir au format Hugging Face ouvre la voie à une compatibilité avec les outils modernes d'inférence et de fine-tuning.

Le processus consiste à extraire les poids depuis mp_rank_00_model_states.pt, à les mapper vers la structure GPT2Model, puis à créer des classes personnalisées pour transformers.

Structure des poids : Megatron-LM vs Hugging Face

Les poids originaux utilisent une projection fusionnée QKV [6144, 2048], qui doit être divisée en matrices séparées Q, K, V de taille [2048, 2048]. Cartographie complète :

| Megatron-LM | Hugging Face |

|-------------|--------------|

| word_embeddings.weight | model.embed_tokens.weight |

| position_embeddings.weight | model.embed_positions.weight |

| transformer.layers.{i}.input_layernorm. | model.layers.{i}.input_layernorm. |

| transformer.layers.{i}.attention.query_key_value.weight | model.layers.{i}.self_attn.{q,k,v}_proj.weight |

| transformer.layers.{i}.attention.query_key_value.bias | model.layers.{i}.self_attn.{q,k,v}_proj.bias |

| transformer.layers.{i}.attention.dense. | model.layers.{i}.self_attn.o_proj. |

| transformer.layers.{i}.post_attention_layernorm. | model.layers.{i}.post_attention_layernorm. |

| transformer.layers.{i}.mlp.dense_h_to_4h. | model.layers.{i}.mlp.up_proj. |

| transformer.layers.{i}.mlp.dense_4h_to_h. | model.layers.{i}.mlp.down_proj. |

| transformer.final_layernorm. | model.norm. |

| - | lm_head.weight (copie de embed_tokens) |

Le script convert.py transforme le checkpoint en safetensors compatible avec transformers.

Classes de modèle personnalisées

Des nouvelles classes ont été développées sans dépendances vers Megatron-LM ou DeepSpeed :

• RuGPT3XLConfig : hérite de PretrainedConfig avec des paramètres (vocab_size=50264, hidden_size=2048, num_layers=24).
• RuGPT3XLAttention : attention multi-têtes avec Q/K/V séparés et DynamicCache.
• RuGPT3XLMLP : bloc MLP (up_proj → GELU → down_proj).
• RuGPT3XLDecoderLayer : couche décodeur (pre-LN → attention → post-LN → MLP).
• RuGPT3XLModel : embeddings + 24 couches + LN finale.
• RuGPT3XLForCausalLM : inclut lm_head.

Améliorations clés :

• Méthode forward() standard pour SFTTrainer/LoRA.
• Cache KV via DynamicCache.
• Checkpointing de gradients.
• Fonctionne sur CPU/GPU avec device_map="auto".

Le modèle est désormais disponible sur Hugging Face sous le nom evilfreelancer/ruGPT3XL.

Résultats des tests

Génération sur RTX 4090 (float16) :

• Vitesse moyenne : 66,7 tokens/sec (batch_size=1).
• Maintien du contexte sur les prompts longs.
• Génération précise (recettes, récits historiques).

Benchmark MERA (score global : 0,198) :

• PARus (sens commun) : 0,500
• ruHateSpeech : 0,558
• BPS (code/math) : 0,528
• RWSD (raisonnement) : 0,488
• ruTiE (dialogue) : 0,502
• ruMMLU : 0,252

Les performances en mathématiques et code sont proches de la base — ce qui est attendu pour un modèle de base sans ajustement par instruction.

Conversion au format GGUF pour llama.cpp

Après la conversion au format HF, le script convert_hf_to_gguf.py a été adapté pour ruGPT-3 XL. Le modèle s’exécute dans llama.cpp sans dépendance à Megatron, avec prise en charge du cache KV et une inférence moderne.

Cela permet la fine-tuning sur des jeux de données personnalisés via LoRA ou SFT, son intégration dans des pipelines, ainsi que son test sur des dispositifs embarqués.

Points clés

• Compatibilité : fonctionne avec transformers ≥4.x, pas besoin de stack hérité.
• Performance : 66,7 t/s sur RTX 4090, support du cache KV pour la génération.
• Apprenabilité : prêt pour LoRA/SFTTrainer avec checkpointing de gradients.
• Benchmarks : fort en sens commun et détection de haine (>0,5), plus faible en mathématiques/code.
• Formats : HF + GGUF pour une large adoption dans l’écosystème.

— Editorial Team