Megatron-LM

Megatron-LM permet d’entraîner de grands modèles de langage transformer à grande échelle. Il offre un parallélisme de modèle efficace basé sur les tenseurs, le pipeline et les séquences pour la pré-entraînement de modèles de langage basés sur les transformers, tels que GPT (Decoder Only), BERT (Encoder Only) et T5 (Encoder-Decoder). Pour des informations détaillées et le fonctionnement en coulisses, veuillez consulter le dépôt github.

Que comprend l'intégration ?

Accelerate intègre la fonctionnalité suivante de Megatron-LM afin de permettre la pré-entraînement/le fine-tuning à grande échelle des modèles BERT (Encoder), GPT (Decoder) ou T5 (Encoder et Decoder) :

a. Parallélisme de tenseurs (TP) : Réduit l'empreinte mémoire sans communication supplémentaire importante sur les rangs intra-nœud. Chaque tenseur est divisé en plusieurs morceaux, chaque fragment résidant sur un GPU distinct. À chaque étape, le même mini-lot de données est traité indépendamment et en parallèle par chaque fragment, suivi d'une synchronisation sur tous les GPU (opération all-reduce). Dans une couche de transformeur simple, cela entraîne 2 all-reduces dans le chemin avant et 2 dans le chemin arrière. Pour plus de détails, veuillez consulter l'article de recherche Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism et cette section du billet de blog The Technology Behind BLOOM Training.

b. Parallélisme de pipeline (PP) : réduit l’empreinte mémoire et permet l’entraînement à grande échelle grâce à la parallélisation inter-nœuds. Réduit la bulle du PP naïf grâce à l’ordonnancement PipeDream-Flush/1F1B et à l’ordonnancement Interleaved 1F1B. Les couches sont réparties uniformément entre les étapes du PP. Par exemple, si un modèle possède 24 couches et que nous disposons de 4 GPU pour le parallélisme de pipeline, chaque GPU aura 6 couches (24/4). Pour plus de détails sur les ordonnancements permettant de réduire le temps d’attente du PP, veuillez consulter l’article de recherche Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM et cette section du billet de blog The Technology Behind BLOOM Training.

c. Parallélisme de séquence (SP) : réduit l’empreinte mémoire sans communication supplémentaire. Applicable uniquement lors de l’utilisation du TP. Il réduit la mémoire d’activation requise en empêchant les mêmes copies d’être présentes sur les rangs du parallélisme de tenseurs après all-reduce en les remplaçant par reduce-scatter et l’opération no-op est remplacée par all-gather. Comme all-reduce = reduce-scatter + all-gather, cela économise une grande quantité de mémoire d’activation sans coût de communication supplémentaire. Pour le dire simplement, il découpe les sorties de chaque couche de transformeur le long de la dimension de séquence, par exemple, si la longueur de séquence est 1024 et la taille du TP est 4, chaque GPU aura 256 jetons (1024/4) pour chaque échantillon. Cela augmente la taille de lot qui peut être prise en charge pour l’entraînement. Pour plus de détails, veuillez consulter l’article de recherche Reducing Activation Recomputation in Large Transformer Models.

d. Parallélisme de données (DP) via Distributed Optimizer : réduit l’empreinte mémoire en répartissant les états de l’optimiseur et les gradients entre les rangs DP (au lieu de la méthode traditionnelle qui consiste à répliquer l’état de l’optimiseur sur les rangs de parallélisme de données). Par exemple, lors de l’utilisation de l’optimiseur Adam avec une formation en précision mixte, chaque paramètre représente 12 octets de mémoire. Cela est réparti égalementment entre les GPU, c’est-à-dire que chaque paramètre représenterait 3 octets (12/4) si nous disposons de 4 GPU. Pour plus de détails, veuillez consulter le papier de recherche ZeRO: Memory Optimizations Toward Training Trillion Parameter Models et la section suivante du blog The Technology Behind BLOOM Training.

e. Parallélisme d'experts (EP) Le parallélisme d'experts dans Megatron-LM est utilisé pour les couches Mixture-of-Experts (MoE), où de nombreux « experts » (petits réseaux feed-forward) existent, mais où seuls quelques-uns sont activés pour chaque jeton. Au lieu de placer tous les experts sur chaque GPU, Megatron distribue différents experts sur différents GPU — c'est le parallélisme d'experts. Pendant l'entraînement, les jetons sont routés vers les GPU qui hébergent leurs experts sélectionnés, y sont calculés, puis renvoyés, ce qui réduit le coût mémoire. Il est souvent combiné avec le parallélisme de tenseurs/pipelines pour les modèles de grande envergure. f. Recalcul complet des activations : Réduit considérablement l'empreinte mémoire des activations grâce à un pointage de contrôle (checkpointing) intelligent des activations. Il ne stocke pas les activations qui occupent une grande quantité de mémoire tout en étant rapides à recalculer, atteignant ainsi un excellent compromis entre mémoire et recalcul. Par exemple, pour GPT-3, cela entraîne une réduction de 70 % de la mémoire requise pour les activations, au prix d'un surcoût de seulement 2,7 % en FLOPs pour le recalcul des activations. Pour plus de détails, veuillez consulter l'article de recherche Reducing Activation Recomputation in Large Transformer Models.

g. Noyaux fusionnés : Softmax fusionné, Layer Norm fusionné en précision mixte et accumulation des gradients fusionnée au calcul des gradients de poids de la couche linéaire. Ajout fusionné GeLU et Bias+Dropout+Residual compilé par le JIT de PyTorch.

h. Prise en charge des jeux de données indexés : format binaire efficace des jeux de données pour l’entraînement à grande échelle. Prise en charge du fichier d’index mmap, cached et du format de chargeur lazy.

i. Redimensionnement des points de contrôle et interopérabilité : Utilitaire permettant de redimensionner les points de contrôle Megatron-LM de tailles de parallélisme tensoriel et de pipeline variables vers les points de contrôle fractionnés de Transformers, qui bénéficient d'un excellent support avec une multitude d'outils tels que Accelerate Big Model Inference, Megatron-DeepSpeed Inference, etc. Le support est également disponible pour convertir les points de contrôle fractionnés de Transformers en points de contrôle Megatron-LM de tailles de parallélisme tensoriel et de pipeline variables pour l'entraînement à grande échelle.

Prérequis

Vous devrez installer les dernières versions de pytorch, cuda, nccl et NVIDIA APEX, ainsi que la bibliothèque nltk. Consultez la documentation pour plus de détails. Une autre façon de configurer l'environnement consiste à récupérer un conteneur NVIDIA PyTorch qui inclut toutes les installations requises depuis NGC.

Voici une méthode étape par étape pour configurer l'environnement conda :

1. Créer un environnement virtuel

conda create --name ml

2. En supposant que la machine a CUDA 11.3 installé, l’installation de la version GPU correspondante de PyTorch

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

3. Installer Nvidia APEX

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
cd ..

4. Installer Megatron-LM

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout 9a1c0d05c992c8a241da384ab27dce2021bb56dd
you need to manually move gpt_builders.py to megatron/training and update
include = [
    "megatron.core", 
    "megatron.core.*",
    "megatron.training",
    "megatron.training.*",
    "megatron.legacy",
    "megatron.legacy.*",
]
in pyproject.toml file to unblock yourself from using Megatron
pip install --no-use-pep517 -e .

Préparer le point de contrôle Megaton-LM

Si vous souhaitez effectuer un fine-tuning d'un modèle, assurez-vous d'avoir un checkpoint au format torch dist prêt. Si vous n'avez accès qu'au modèle huggingface, envisagez de le convertir en un checkpoint au format torch dist acceptable par Megatron. Un exemple consiste à utiliser le script de slime, en prenant les modèles GLM comme exemple :

source /your/path/to/slime/scripts/models/glm4.5-355B-A32B.sh
srun torchrun --nproc-per-node 8 \
   /your/path/to/slime/tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /your/path/to/huggingface/models/GLM4.5-355B-A32B \
    --save /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist

Après la conversion, assurez-vous : 1. sous /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist : modifiez le contenu de latest_checkpointed_iteration.txt de release à 0 et renommez le répertoire release en iter_0000000 ; 2 : dans la configuration, assurez-vous que megatron_lm_no_load_optim est défini sur true afin qu'aucun état d'optimiseur ne soit nécessaire.

Accélérer le plugin Megatron-LM

Les fonctionnalités importantes sont directement prises en charge via la commande accelerate config. Un exemple des questions correspondantes pour l'utilisation des fonctionnalités de Megatron-LM est présenté ci-dessous :

:~$ accelerate config --config_file "megatron_gpt_config.yaml"
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 2
How many different machines will you use (use more than 1 for multi-node training)? [1]: 
Do you want to use DeepSpeed? [yes/NO]: 
Do you want to use FullyShardedDataParallel? [yes/NO]: 
Do you want to use Megatron-LM ? [yes/NO]: yes
What is the Tensor Parallelism degree/size? [1]:2
Do you want to enable Sequence Parallelism? [YES/no]: 
What is the Pipeline Parallelism degree/size? [1]:2
What is the number of micro-batches? [1]:2
Do you want to enable selective activation recomputation? [YES/no]: 
Do you want to use distributed optimizer which shards optimizer state and gradients across data parallel ranks? [YES/no]: 
What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
How many GPU(s) should be used for distributed training? [1]:4
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16

La configuration résultante est présentée ci-dessous :

~$ cat megatron_gpt_config.yaml 
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MEGATRON_LM
downcast_bf16: 'no'
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config:
  megatron_lm_gradient_clipping: 1.0
  megatron_lm_num_micro_batches: 2
  megatron_lm_pp_degree: 2
  megatron_lm_recompute_activations: true
  megatron_lm_sequence_parallelism: true
  megatron_lm_tp_degree: 2
  megatron_lm_use_distributed_optimizer: true
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false

Nous prendrons l’exemple de la pré-formation de GPT. Les modifications minimales requises dans le run_clm_no_trainer.py officiel pour utiliser Megatron-LM sont les suivantes :

1. Megatron-LM utilisant sa propre implémentation d'Optimizer, il est nécessaire d'utiliser le scheduler correspondant compatible avec celui-ci. Ainsi, seul le scheduler de Megatron-LM est pris en charge. L'utilisateur devra créer accelerate.utils.MegatronLMDummyScheduler. Un exemple est donné ci-dessous :

from accelerate.utils import MegatronLMDummyScheduler

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    lr_scheduler = MegatronLMDummyScheduler(
        optimizer=optimizer,
        total_num_steps=args.max_train_steps,
        warmup_num_steps=args.num_warmup_steps,
    )
else:
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
    )

2. Obtenir les détails de la taille totale du lot nécessite désormais de prendre en compte les tailles du parallélisme de tenseurs et du parallélisme de pipeline. Un exemple de récupération de la taille totale effective du lot est présenté ci-dessous :

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
else:
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

3. Lors de l'utilisation de Megatron-LM, les pertes sont déjà moyennées sur le groupe de parallélisme des données

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses.append(loss)
else:
    losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses = torch.tensor(losses)
else:
    losses = torch.cat(losses)

4. Pour Megatron-LM, nous devons enregistrer le modèle à l'aide de accelerator.save_state

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    accelerator.save_state(args.output_dir)
else:
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(
        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
    )

C’est tout ! Nous sommes prêts à démarrer 🚀. Vous trouverez le script d’exemple dans le dossier examples, à l’emplacement accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py. Exécutons-le pour l’architecture de modèle gpt-large en utilisant 4 GPU A100-80GB.

accelerate launch --config_file megatron_gpt_config.yaml \
examples/by_feature/megatron_lm_gpt_pretraining.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--block_size 1024 \
--learning_rate 5e-5 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 5 \
--with_tracking \
--report_to "wandb" \
--output_dir "awesome_model"

Voici quelques extraits importants des journaux de sortie :

Loading extension module fused_dense_cuda...
>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds
 > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)
Building gpt model in the pre-training mode.
The Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.
Preparing dataloader
Preparing dataloader
Preparing model
 > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120
 > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120
Preparing optimizer
Preparing scheduler
> learning rate decay style: linear
10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****
10/10/2022 22:57:22 - INFO - __main__ -   Num examples = 2318
10/10/2022 22:57:22 - INFO - __main__ -   Num Epochs = 5
10/10/2022 22:57:22 - INFO - __main__ -   Instantaneous batch size per device = 24
10/10/2022 22:57:22 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 48
10/10/2022 22:57:22 - INFO - __main__ -   Gradient Accumulation steps = 1
10/10/2022 22:57:22 - INFO - __main__ -   Total optimization steps = 245
 20%|████████████▍                                                 | 49/245 [01:04<04:09,  1.27s/it]
 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209
 40%|████████████████████████▊                                     | 98/245 [02:10<03:07,  1.28s/it]
 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359
 60%|████████████████████████████████████▌                        | 147/245 [03:16<02:05,  1.28s/it]
 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344
 80%|████████████████████████████████████████████████▊            | 196/245 [04:22<01:02,  1.28s/it]
 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874
100%|█████████████████████████████████████████████████████████████| 245/245 [05:27<00:00,  1.28s/it]

Il existe un grand nombre d'autres options/paramètres que l'on peut définir à l'aide de accelerate.utils.MegatronLMPlugin.

Fonctionnalités avancées pour écrire une étape d'entraînement personnalisée et utiliser les jeux de données indexés de Megatron-LM

Pour tirer parti de davantage de fonctionnalités, veuillez consulter les détails ci-dessous.

1. Voici un exemple des modifications nécessaires pour personnaliser le Train Step lors de l’utilisation de Megatron-LM. Vous implémenterez accelerate.utils.AbstractTrainStep ou hériterez de leurs enfants correspondants accelerate.utils.GPTTrainStep, accelerate.utils.BertTrainStep ou accelerate.utils.T5TrainStep.

from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group


# Custom loss function for the Megatron model
class GPTTrainStepWithCustomLoss(GPTTrainStep):
    def __init__(self, megatron_args, **kwargs):
        super().__init__(megatron_args)
        self.kwargs = kwargs

    def get_loss_func(self):
        def loss_func(inputs, loss_mask, output_tensor):
            batch_size, seq_length = output_tensor.shape
            losses = output_tensor.float()
            loss_mask = loss_mask.view(-1).float()
            loss = losses.view(-1) * loss_mask

            # Resize and average loss per sample
            loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)
            loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)
            loss_per_sample = loss_per_sample / loss_mask_per_sample

            # Calculate and scale weighting
            weights = torch.stack([(inputs == kt).float() for kt in self.kwargs["keytoken_ids"]]).sum(axis=[0, 2])
            weights = 1.0 + self.kwargs["alpha"] * weights
            # Calculate weighted average
            weighted_loss = (loss_per_sample * weights).mean()

            # Reduce loss across data parallel groups
            averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])

            return weighted_loss, {"lm loss": averaged_loss[0]}

        return loss_func

    def get_forward_step_func(self):
        def forward_step(data_iterator, model):
            """Forward step."""
            # Get the batch.
            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)

            return output_tensor, partial(self.loss_func, tokens, loss_mask)

        return forward_step


def main():
    # Custom loss function for the Megatron model
    keytoken_ids = []
    keywords = ["plt", "pd", "sk", "fit", "predict", " plt", " pd", " sk", " fit", " predict"]
    for keyword in keywords:
        ids = tokenizer([keyword]).input_ids[0]
        if len(ids) == 1:
            keytoken_ids.append(ids[0])
    accelerator.print(f"Keytoken ids: {keytoken_ids}")
    accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss
    accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {
        "keytoken_ids": keytoken_ids,
        "alpha": 0.25,
    }

2. Pour utiliser les jeux de données Megatron-LM, quelques modifications supplémentaires sont nécessaires. Les chargeurs de données pour ces jeux de données ne sont disponibles que sur le rang 0 de chaque groupe de parallélisme de tenseurs. Par conséquent, il existe des rangs où le chargeur de données ne sera pas disponible, ce qui nécessite des ajustements dans la boucle d'entraînement. Le fait de pouvoir effectuer toutes ces opérations montre à quel point Accelerate est flexible et extensible. Les modifications requises sont les suivantes.

a. Pour les jeux de données indexés Megatron-LM, il est nécessaire d'utiliser MegatronLMDummyDataLoader et de lui passer les arguments de jeu de données requis, tels que data_path, seq_length, etc. Voir ici pour la liste des arguments disponibles.

from accelerate.utils import MegatronLMDummyDataLoader

megatron_dataloader_config = {
    "data_path": args.data_path,
    "splits_string": args.splits_string,
    "seq_length": args.block_size,
    "micro_batch_size": args.per_device_train_batch_size,
}
megatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)
accelerator.state.megatron_lm_plugin.megatron_dataset_flag = True

b. megatron_dataloader est répété 3 fois pour obtenir les chargeurs de données d'entraînement, de validation et de test, conformément aux proportions définies par args.splits_string

model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(
    model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader
)

c. Modifications des boucles d’entraînement et d’évaluation, le dataloader n’étant disponible que sur les rangs 0 du parallélisme de tenseurs Il faut donc itérer uniquement si le dataloader n’est pas None, sinon fournir un dictionnaire vide Ainsi, on boucle à l’aide d’une boucle while et on interrompt lorsque completed_steps est égal à args.max_train_steps. Cela ressemble à la configuration de Megatron-LM, où l’utilisateur doit fournir max_train_steps lors de l’utilisation de jeux de données indexés par Megaton-LM. Cela montre à quel point Accelerate est flexible et extensible.

while completed_steps < args.max_train_steps:
    model.train()
    batch = next(train_dataloader) if train_dataloader is not None else {}
    outputs = model(**batch)
    loss = outputs.loss
    ...

    if completed_steps % eval_interval == 0:
        eval_completed_steps = 0
        losses = []
        while eval_completed_steps < eval_iters:
            model.eval()
            with torch.no_grad():
                batch = next(eval_dataloader) if eval_dataloader is not None else {}
                outputs = model(**batch)

Utilitaire pour le remodelage de points de contrôle et l'interopérabilité

1. Les scripts correspondants sont présents dans la bibliothèque Transformers, sous les modèles respectifs. Actuellement, il est disponible pour le modèle GPT checkpoint_reshaping_and_interoperability.py

2. Voici un exemple de conversion d’un point de contrôle de Megatron-LM vers un point de contrôle universel Transformers sharded.

python checkpoint_reshaping_and_interoperability.py \
--convert_checkpoint_from_megatron_to_transformers \
--load_path "gpt/iter_0005000" \
--save_path "gpt/trfs_checkpoint" \
--max_shard_size "200MB" \
--tokenizer_name "gpt2" \
--print-checkpoint-structure

3. Conversion d'un checkpoint de transformers vers megatron avec tp_size=2, pp_size=2 et dp_size=2.

python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \
--load_path "gpt/trfs_checkpoint" \
--save_path "gpt/megatron_lm_checkpoint" \
--target_tensor_model_parallel_size 2 \
--target_pipeline_model_parallel_size 2 \
--target_data_parallel_size 2 \
--target_params_dtype "bf16" \
--make_vocab_size_divisible_by 128 \
--use_distributed_optimizer \
--print-checkpoint-structure

Les modèles GPT de Megatron-LM prennent en charge le retour des logits et la fonction megatron_generate pour la génération de texte

1. Renvoyer les logits nécessite de définir require_logits=True dans MegatronLMPlugin, comme indiqué ci-dessous. Ces derniers seront disponibles dans la dernière étape du pipeline.

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)

2. Méthode megatron_generate pour le modèle GPT de Megatron-LM : elle utilise le parallélisme de tenseurs et le parallélisme de pipeline pour effectuer la génération d’un lot d’entrées lors de l’utilisation de la génération gloutonne avec ou sans échantillonnage top_k/top_p, et pour des entrées de prompt individuelles lors de l’utilisation du décodage par recherche de faisceau. Seul un sous-ensemble des fonctionnalités de generate de transformers est pris en charge. Cela permet d’utiliser de grands modèles via le parallélisme de tenseurs et de pipeline pour la génération (le cache des clés-valeurs est déjà effectué et les noyaux fusionnés sont utilisés par défaut). Cela exige que la taille du parallélisme de données soit de 1, et que le parallélisme de séquences et le point de contrôle des activations soient désactivés. Il est également nécessaire de spécifier le chemin du fichier de vocabulaire et du fichier de fusions du tokeniseur. L’exemple ci-dessous montre comment configurer et utiliser la méthode megatron_generate pour le modèle GPT de Megatron-LM.

# specifying tokenizer's vocab and merges file
vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

# inference using `megatron_generate` functionality
tokenizer.pad_token = tokenizer.eos_token
max_new_tokens = 64
batch_texts = [
    "Are you human?",
    "The purpose of life is",
    "The arsenal was constructed at the request of",
    "How are you doing these days?",
]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)

# top-p sampling
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_p=0.8,
    top_p_decay=0.5,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# top-k sampling
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_k=50,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# adding `bos` token at the start
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"], batch_encodings["attention_mask"], max_new_tokens=max_new_tokens, add_BOS=True
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# beam search => only takes single prompt
batch_texts = ["The purpose of life is"]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    num_beams=20,
    length_penalty=1.5,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

3. Un exemple de bout en bout de l'utilisation de la méthode megatron_generate pour le modèle GPT de Megatron-LM est disponible dans megatron_gpt2_generation.py avec le fichier de configuration megatron_lm_gpt_generate_config.yaml. Le script bash avec la commande accelerate launch est disponible dans megatron_lm_gpt_generate.sh. Les journaux de sortie du script sont disponibles dans megatron_lm_gpt_generate.log.

Prise en charge des embeddings positionnels ROPE et ALiBi et de l'attention Multi-Query

1. Pour l’attention ROPE/ALiBi, passez position_embedding_type avec ("absolute" | "rotary" | "alibi") à MegatronLMPlugin comme indiqué ci-dessous.

other_megatron_args = {"position_embedding_type": "alibi"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

2. Pour l’attention multi-requête, passez attention_head_type avec ("multihead" | "multiquery") à MegatronLMPlugin comme indiqué ci-dessous.

other_megatron_args = {"attention_head_type": "multiquery"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

Limites

1. Prend en charge les modèles GPT2, Megatron-BERT et T5 de Transformers. Cela couvre les classes de modèles Decoder only, Encode only et Encoder-Decoder.

2. Seule la perte est retournée par le passage avant du modèle, car il existe une interaction complexe entre le parallélisme de pipeline, de tenseurs et de données en arrière-plan. L'appel model(**batch_data) renvoie la ou les pertes moyennées sur les rangs de parallélisme de données. C'est suffisant dans la plupart des cas où les tâches de pré-entraînement sont exécutées en utilisant les fonctionnalités de Megatron-LM et où vous pouvez facilement calculer la perplexity à partir de la perte. Pour le modèle GPT, le retour des logits en plus de la ou des pertes est pris en charge. Ces logits ne sont pas regroupés sur les rangs de parallélisme de données. Utilisez accelerator.utils.gather_across_data_parallel_groups pour regrouper les logits sur les rangs de parallélisme de données. Ces logits, ainsi que les étiquettes, peuvent être utilisés pour calculer diverses métriques de performance.

3. Le processus principal est le dernier rang, car les pertes/logits sont disponibles dans la dernière étape du pipeline. accelerator.is_main_process et accelerator.is_local_main_process renvoient True pour le dernier rang lors de l'utilisation de l'intégration Megatron-LM.

4. Lors de l'appel à accelerator.prepare, un modèle Megatron-LM correspondant au modèle Transformers donné est créé avec des poids aléatoires. Veuillez utiliser accelerator.load_state pour charger le checkpoint Megatron-LM avec les partitions TP, PP et DP correspondantes.

5. Actuellement, le redimensionnement des points de contrôle et la prise en charge de l'interopérabilité ne sont disponibles que pour GPT. Bientôt, ils seront étendus à BERT et T5.

6. gradient_accumulation_steps doit être égal à 1. Lors de l’utilisation de Megatron-LM, les micro-lots dans le parallélisme de pipeline sont synonymes d’accumulation de gradients.

7. Lors de l'utilisation de Megatron-LM, utilisez accelerator.save_state et accelerator.load_state pour enregistrer et charger les points de contrôle.

8. Voici la correspondance entre les architectures de modèles Megatron-LM et les architectures de modèles transformers équivalentes. Seules ces architectures de modèles transformers sont prises en charge.

a. Megatron-LM BertModel : modèles transformers dont le type de modèle dans la configuration contient megatron-bert, par exemple MegatronBERT

b. Megatron-LM GPTModel : modèles transformers dont le type de modèle dans la configuration contient gpt2, par exemple OpenAI GPT2

c. Megatron-LM T5Model : modèles transformers dont le type de modèle dans la configuration contient t5, par exemple T5 et MT5