[megatron] support megatron num_train_epochs (#4432)

Author: Jintao-Huang

docs/source/Instruction/Megatron-SWIFT训练.md

@@ -300,3 +300,4 @@ Megatron训练参数继承自Megatron参数和基本参数。基本参数的内
 - 🔥streaming: 流式读取并处理数据集，默认False。通常在处理大型数据集时，设置为True。更多流式的参数查看命令行参数文档。
 - lazy_tokenize: 默认为False。若该参数设置为False，则在训练之前对所有的数据集样本进行tokenize（这可以避免在训练中出现报错）；设置为True，则在训练中对数据集进行tokenize（这可以节约内存）。
 - max_epochs: 训练到`max_epochs`时强制退出训练，并对权重进行验证和保存。该参数在使用流式数据集时很有用。默认为None。
+  - 注意：如果你使用非流式数据集，该参数会为你自动计算train_iters，你不需要手动传入`train_iters`。

docs/source_en/Instruction/Megatron-SWIFT-Training.md

@@ -311,3 +311,4 @@ Megatron training parameters inherit from Megatron parameters and basic paramete
 - 🔥streaming: Stream reading and processing of the dataset, default is False. It is typically set to True when handling large datasets. For more information on streaming parameters, refer to the command-line parameters documentation.
 - lazy_tokenize: Default is False. If this parameter is set to False, all dataset samples are tokenized before training (this avoids errors during training); if set to True, tokenization occurs during training (this saves memory).
 - max_epochs: Forces the training to exit after reaching `max_epochs`, and performs validation and saving of the model weights. This parameter is especially useful when using a streaming dataset. Default is None.
+  - Note: If you use a non-streaming dataset, this parameter will automatically calculate train_iters for you, so there is no need to pass `train_iters` manually.

swift/cli/_megatron/main.py

@@ -7,8 +7,8 @@
 logger = get_logger()
 
 ROUTE_MAPPING: Dict[str, str] = {
-    'sft': 'swift.cli._megatron.sft',
     'pt': 'swift.cli._megatron.pt',
+    'sft': 'swift.cli._megatron.sft',
 }
 
 

swift/llm/model/patcher.py

@@ -350,7 +350,8 @@ def new_get_cached_module_file(pretrained_model_name_or_path, *args, **kwargs):
 
 @contextmanager
 def patch_tp_plan(load_model: bool):
-    if not load_model or not is_mp_ddp() or version.parse(transformers.__version__) < version.parse('4.50'):
+    if not load_model or not is_mp_ddp() or version.parse(
+            transformers.__version__) < version.parse('4.50') or 'WORLD_SIZE' not in os.environ:
         yield
         return
     WORLD_SIZE = os.environ.get('WORLD_SIZE')

swift/megatron/init.py

@@ -1,7 +1,6 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
 import sys
-from contextlib import contextmanager
 
 from swift.llm import git_clone_github
 from swift.utils import get_logger, is_megatron_available, safe_ddp_context, subprocess_run
@@ -30,50 +29,6 @@ def _patch_transformer_engine():
             pass
 
 
-def new_cyclic_iter(iter):
-    from megatron.training import get_args
-    args = get_args()
-    max_epochs = args.max_epochs
-    i = 0
-    while True:
-        if getattr(args, 'is_training', False):
-            if max_epochs and i >= max_epochs:
-                logger.info(f'Training of {i} epochs has been completed, the training has finished.')
-                break
-            logger.info(f'The training of Epoch {i} starts...')
-        for x in iter:
-            yield x
-        i += 1
-
-
-@contextmanager
-def _training_context():
-    from megatron.training import get_args
-    args = get_args()
-    args.is_training = True
-    try:
-        yield
-    finally:
-        args.is_training = False
-
-
-def _patch_max_epochs():
-    # support max_epochs
-    from megatron.training import training
-    train_step_origin = training.train_step
-
-    def train_step(*args, **kwargs):
-        with _training_context():
-            try:
-                return train_step_origin(*args, **kwargs)
-            except StopIteration:
-                return {}, True, True, True, 0, None, None
-
-    training.train_step = train_step
-
-    training.cyclic_iter = new_cyclic_iter
-
-
 def _patch__batched_p2p_ops():
     from megatron.core.pipeline_parallel import p2p_communication
 
@@ -88,7 +43,6 @@ def _batched_p2p_ops(**kwargs):
 
 def _patch_megatron():
     _patch_transformer_engine()
-    _patch_max_epochs()
     _patch__batched_p2p_ops()
 
 

swift/megatron/train/__init__.py

@@ -1,2 +1,3 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
 from .pt import megatron_pt_main
 from .sft import megatron_sft_main

swift/megatron/train/sft.py

@@ -1,19 +1,25 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
+from contextlib import contextmanager
+from functools import partial
 from typing import List, Union
 
+from megatron.core import mpu
 from megatron.core.enums import ModelType
-from megatron.training import pretrain
+from megatron.core.utils import StragglerDetector
+from megatron.training import get_args, get_timers, pretrain, training
 
 from swift.llm.train import SwiftSft
 from swift.utils import get_logger, is_master, plot_images
 from ..argument import MegatronTrainArguments
 from ..utils import patch_megatron_tokenizer
 from .patcher import patch_megatron_data_collator
-from .utils import build_streaming_dataloader, forward_step, get_swift_datasets_provider
+from .utils import build_streaming_dataloader, get_batch, get_swift_datasets_provider
 
 logger = get_logger()
 
+stimer = StragglerDetector()
+
 
 class MegatronSft(SwiftSft):
     args_class = MegatronTrainArguments
@@ -30,8 +36,92 @@ def __init__(self, args: Union[List[str], MegatronTrainArguments, None] = None)
         self.template.use_megatron = True
         args.save_args(args.save)
 
+    @contextmanager
+    def _get_train_iters(self, train_dataset):
+        from megatron.training import training
+        origin_initialize_megatron = training.initialize_megatron
+
+        def initialize_megatron(*_args, **kwargs):
+            res = origin_initialize_megatron(*_args, **kwargs)
+            args = get_args()
+            if args.train_iters is None and hasattr(train_dataset, '__len__'):
+                data_parallel_size = mpu.get_data_parallel_world_size()
+                step_batch_size = \
+                    args.micro_batch_size * data_parallel_size
+                dataset_sample = len(train_dataset) // step_batch_size * step_batch_size
+                args.train_iters = (dataset_sample * args.max_epochs // args.global_batch_size) + 1
+            return res
+
+        training.initialize_megatron = initialize_megatron
+        try:
+            yield
+        finally:
+            training.initialize_megatron = origin_initialize_megatron
+
+    @staticmethod
+    def new_cyclic_iter(iter):
+        args = get_args()
+        max_epochs = args.max_epochs
+        i = 0
+        while True:
+            if getattr(args, 'is_training', False):
+                if max_epochs and i >= max_epochs:
+                    logger.info(f'Training of {i} epochs has been completed, the training has finished.')
+                    break
+                logger.info(f'The training of Epoch {i} starts...')
+            for x in iter:
+                yield x
+            i += 1
+
+    @staticmethod
+    @contextmanager
+    def _training_context():
+        args = get_args()
+        args.is_training = True
+        try:
+            yield
+        finally:
+            args.is_training = False
+
+    def train_step(self, forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config):
+        return self._train_step_origin(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config)
+
+    def _patch_train_step(self):
+        # support max_epochs
+        def train_step(*args, **kwargs):
+            with self._training_context():
+                try:
+                    return self.train_step(*args, **kwargs)
+                except StopIteration:
+                    return {}, True, True, True, 0, None, None
+
+        self._train_step_origin = training.train_step
+        training.train_step = train_step
+        training.cyclic_iter = MegatronSft.new_cyclic_iter
+
+    def forward_step(self, data_iterator, model):
+        from pretrain_gpt import loss_func
+
+        timers = get_timers()
+
+        # Get the batch.
+        timers('batch-generator', log_level=2).start()
+        global stimer
+        with stimer(bdata=True):
+            data = get_batch(data_iterator)
+        if not data:
+            raise StopIteration
+        timers('batch-generator').stop()
+
+        with stimer:
+            output_tensor = model(**data)
+        labels = data.get('labels')
+        loss_mask = None if labels is None else (labels != -100).float()
+        return output_tensor, partial(loss_func, loss_mask)
+
     def run(self):
         args = self.args
+        self._patch_train_step()
 
         train_dataset, val_dataset = self._get_dataset()
         train_dataset, val_dataset = self._encode_dataset(train_dataset, val_dataset)
@@ -46,13 +136,13 @@ def run(self):
         logging_path = os.path.join(args.save, 'logging.jsonl')
         logger.info(f'The logging file will be saved in: {logging_path}')
         try:
-            with patch_megatron_data_collator(data_collator):
+            with patch_megatron_data_collator(data_collator), self._get_train_iters(train_dataset):
                 extra_args_provider = args.megatron_model_meta.extra_args_provider
                 pretrain(
                     datasets_provider,
                     args.megatron_model_meta.model_provider,
                     ModelType.encoder_or_decoder,
-                    forward_step,
+                    self.forward_step,
                     extra_args_provider=extra_args_provider,
                     args_defaults=args.extra_args)
         finally:

swift/megatron/train/utils.py

@@ -1,18 +1,14 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from functools import partial
 from typing import Any, Dict, Optional
 
 import torch
 from megatron.core import mpu
 from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.utils import StragglerDetector
 from megatron.training import get_args, get_timers
 from megatron.training.training import cyclic_iter
 
 from swift.llm import DataLoaderDispatcher
 
-stimer = StragglerDetector()
-
 
 def get_swift_datasets_provider(train_dataset, val_dataset):
 
@@ -67,10 +63,10 @@ def _broadcast(item):
         except StopIteration:
             seq_length = -1
         else:
-            tokens = data['input_ids']
-            seq_length = tokens.shape[1]
+            input_ids = data['input_ids']
+            seq_length = input_ids.shape[1]
             batch = {
-                'tokens': tokens.cuda(non_blocking=True),
+                'input_ids': input_ids.cuda(non_blocking=True),
                 'labels': data['labels'].cuda(non_blocking=True),
                 'attention_mask':
                 None if 'attention_mask' not in data else data['attention_mask'].cuda(non_blocking=True),
@@ -81,13 +77,13 @@ def _broadcast(item):
         if seq_length.item() == -1:
             return {}
         if args.pipeline_model_parallel_size == 1:
-            _broadcast(batch['tokens'])
+            _broadcast(batch['input_ids'])
             _broadcast(batch['labels'])
             _broadcast(batch['attention_mask'])
             _broadcast(batch['position_ids'])
 
         elif mpu.is_pipeline_first_stage():
-            _broadcast(batch['tokens'])
+            _broadcast(batch['input_ids'])
             _broadcast(batch['attention_mask'])
             _broadcast(batch['position_ids'])
 
@@ -102,7 +98,7 @@ def _broadcast(item):
         if seq_length.item() == -1:
             return {}
         micro_batch_size = 1  # use qkv_format 'thd'
-        tokens = torch.empty((micro_batch_size, seq_length), dtype=torch.int64, device=torch.cuda.current_device())
+        input_ids = torch.empty((micro_batch_size, seq_length), dtype=torch.int64, device=torch.cuda.current_device())
         labels = torch.empty((micro_batch_size, seq_length), dtype=torch.int64, device=torch.cuda.current_device())
         if args.create_attention_mask_in_dataloader:
             attention_mask = torch.empty((micro_batch_size, 1, seq_length, seq_length),
@@ -115,26 +111,31 @@ def _broadcast(item):
                                    device=torch.cuda.current_device())
 
         if args.pipeline_model_parallel_size == 1:
-            _broadcast(tokens)
+            _broadcast(input_ids)
             _broadcast(labels)
             _broadcast(attention_mask)
             _broadcast(position_ids)
 
         elif mpu.is_pipeline_first_stage():
             labels = None
 
-            _broadcast(tokens)
+            _broadcast(input_ids)
             _broadcast(attention_mask)
             _broadcast(position_ids)
 
         elif mpu.is_pipeline_last_stage():
-            tokens = None
+            input_ids = None
 
             _broadcast(labels)
             _broadcast(attention_mask)
             _broadcast(position_ids)  # compat packing & cp
 
-        batch = {'tokens': tokens, 'labels': labels, 'attention_mask': attention_mask, 'position_ids': position_ids}
+        batch = {
+            'input_ids': input_ids,
+            'labels': labels,
+            'attention_mask': attention_mask,
+            'position_ids': position_ids
+        }
 
     return batch
 
@@ -213,25 +214,4 @@ def get_batch(data_iterator):
     batch['packed_seq_params'] = get_packed_seq_params(batch['position_ids'])
     # slice batch along sequence dimension for context parallelism
     batch = get_batch_on_this_cp_rank(batch)
-    return batch.values()
-
-
-def forward_step(data_iterator, model):
-    from pretrain_gpt import loss_func
-
-    timers = get_timers()
-
-    # Get the batch.
-    timers('batch-generator', log_level=2).start()
-    global stimer
-    with stimer(bdata=True):
-        data = get_batch(data_iterator)
-    if not data:
-        raise StopIteration
-    tokens, labels, attention_mask, position_ids, packed_seq_params = data
-    timers('batch-generator').stop()
-
-    with stimer:
-        output_tensor = model(tokens, position_ids, attention_mask, labels=labels, packed_seq_params=packed_seq_params)
-    loss_mask = None if labels is None else (labels != -100).float()
-    return output_tensor, partial(loss_func, loss_mask)
+    return batch