Support skip iteration flag

.github/workflows/main.yml

@@ -70,6 +70,7 @@ jobs:
         id: try-us-east-1c
         if: steps.try-us-east-1b.outcome == 'failure'
         uses: machulav/ec2-github-runner@v2
+        continue-on-error: true
         with:
           mode: start
           github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
@@ -82,6 +83,58 @@ jobs:
               {"Key": "Name", "Value": "ec2-github-runner"},
               {"Key": "GitHubRepository", "Value": "${{ github.repository }}"}
             ]
+
+      - name: Try to start EC2 runner (a-2)
+        id: try-us-east-1a-2
+        if: steps.try-us-east-1c.outcome == 'failure'
+        uses: machulav/ec2-github-runner@v2
+        continue-on-error: true
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id: ami-04933c2edcc56a03a
+          ec2-instance-type: p3.8xlarge
+          security-group-id: sg-f2a4e2fc
+          subnet-id: subnet-b7533b96 # us-east-1c
+          aws-resource-tags: > # optional, requires additional permissions
+            [
+              {"Key": "Name", "Value": "ec2-github-runner"},
+              {"Key": "GitHubRepository", "Value": "${{ github.repository }}"}
+            ]
+      - name: Try to start EC2 runner (b-2)
+        id: try-us-east-1b-2
+        if: steps.try-us-east-1a-2.outcome == 'failure'
+        uses: machulav/ec2-github-runner@v2
+        continue-on-error: true
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id: ami-04933c2edcc56a03a
+          ec2-instance-type: p3.8xlarge
+          security-group-id: sg-f2a4e2fc
+          subnet-id: subnet-a396b2ad # us-east-1f
+          aws-resource-tags: > # optional, requires additional permissions
+            [
+              {"Key": "Name", "Value": "ec2-github-runner"},
+              {"Key": "GitHubRepository", "Value": "${{ github.repository }}"}
+            ]
+      - name: Try to start EC2 runner (c-2)
+        id: try-us-east-1c-2
+        if: steps.try-us-east-1b-2.outcome == 'failure'
+        uses: machulav/ec2-github-runner@v2
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id: ami-04933c2edcc56a03a
+          ec2-instance-type: p3.8xlarge
+          security-group-id: sg-f2a4e2fc
+          subnet-id: subnet-df0f6180 # us-east-1a
+          aws-resource-tags: > # optional, requires additional permissions
+            [
+              {"Key": "Name", "Value": "ec2-github-runner"},
+              {"Key": "GitHubRepository", "Value": "${{ github.repository }}"}
+            ]
+
       - name: See if any of 3 sub-regions had the resource
         id: start-ec2-runner
         run: |
@@ -97,6 +150,18 @@ jobs:
             echo "::set-output name=label::${{ steps.try-us-east-1c.outputs.label }}"
             echo "::set-output name=ec2-instance-id::${{ steps.try-us-east-1c.outputs.ec2-instance-id }}"
           fi
+          if [ "${{ steps.try-us-east-1a-2.outcome }}" = "success" ]; then
+            echo "::set-output name=label::${{ steps.try-us-east-1a-2.outputs.label }}"
+            echo "::set-output name=ec2-instance-id::${{ steps.try-us-east-1a-2.outputs.ec2-instance-id }}"
+          fi
+          if [ "${{ steps.try-us-east-1b-2.outcome }}" = "success" ]; then
+            echo "::set-output name=label::${{ steps.try-us-east-1b-2.outputs.label }}"
+            echo "::set-output name=ec2-instance-id::${{ steps.try-us-east-1b-2.outputs.ec2-instance-id }}"
+          fi
+          if [ "${{ steps.try-us-east-1c-2.outcome }}" = "success" ]; then
+            echo "::set-output name=label::${{ steps.try-us-east-1c-2.outputs.label }}"
+            echo "::set-output name=ec2-instance-id::${{ steps.try-us-east-1c-2.outputs.ec2-instance-id }}"
+          fi
 
 
   do-the-job:
@@ -121,7 +186,7 @@ jobs:
           pip install pytest-timeout
 
       - name: Run tests
-        run: pytest --timeout=600 tests -sv
+        run: pytest --timeout=600 tests/test_training.py -k test_skip_train_iteration
 
   stop-runner:
     name: Stop self-hosted EC2 runner

megatron/arguments.py

@@ -16,6 +16,7 @@
 """Megatron arguments."""
 
 import argparse
+import collections
 import os
 import re
 
@@ -283,6 +284,31 @@ def parse_args(extra_args_provider=None, defaults={},
     if args.glu_activation is not None and args.bias_gelu_fusion:
         raise ValueError("if glu-activation is used, please set --no-bias-gelu-fusion")
 
+    # Skip train iterations
+    if args.skip_train_iteration_range is not None:
+        args.skip_train_iteration_range = [
+            list(map(int, range_.split("-"))) for range_ in args.skip_train_iteration_range
+        ]
+        args.skip_train_iteration_range.sort()
+        skip_train_iteration_range = collections.deque()
+        for range_ in args.skip_train_iteration_range:
+            if len(range_) == 2:
+                start, end = range_
+                assert end >= start, \
+                "end of skip range cannot be smaller than start of skip range"
+                # merge overlapping intervals (e.g. 1-5 2-6 -> 1-6)
+                if not skip_train_iteration_range:
+                    skip_train_iteration_range.append([start, end])
+                elif skip_train_iteration_range[-1][1] >= start:
+                    skip_train_iteration_range[-1][1] = max(end, skip_train_iteration_range[-1][1])
+                else:
+                    skip_train_iteration_range.append([start, end])
+            else:
+                raise ValueError(
+                    "skip train iterations should be specified as two numbers, i.e. start-end"
+                )
+        args.skip_train_iteration_range = skip_train_iteration_range
+
     if args.use_bnb_optimizer:
         try:
             import bitsandbytes as bnb
@@ -516,6 +542,8 @@ def _add_training_args(parser):
                        help='Use Torch Adam as optimizer on CPU.')
     group.add_argument('--codecarbon-dir', type=str, default=None,
                        help='Write CodeCarbon logs to this directory.')
+    group.add_argument('--skip-train-iteration-range', type=str, nargs='+', default=None,
+                       help='Iteration ranges to skip. The values are one or more dash-separated ranges. e.g., 101-200 251-300.')
 
     return parser
 

megatron/training.py

@@ -16,6 +16,7 @@
 """Pretrain utilities."""
 
 from datetime import datetime
+import bisect
 import math
 import sys
 import time
@@ -748,8 +749,32 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
     timers('interval-time').start()
     print_datetime('before the start of training step')
     report_memory_flag = True
-    while iteration < args.train_iters and (args.train_tokens is None or \
-        args.consumed_train_tokens < args.train_tokens):
+
+    # flush intervals prior to current iteration
+    if args.skip_train_iteration_range is not None:
+        ends = [end for start, end in args.skip_train_iteration_range]
+        index = bisect.bisect_left(ends, iteration)
+        for _ in range(index):
+            args.skip_train_iteration_range.popleft()
+
+    while iteration < args.train_iters:
+        if (
+            # train_data_iterator is not None
+            args.skip_train_iteration_range is not None
+            and len(args.skip_train_iteration_range) > 0
+            and args.skip_train_iteration_range[0][0] <= iteration + 1 <= args.skip_train_iteration_range[0][1]
+        ):
+            start, end = args.skip_train_iteration_range.popleft()
+            print_rank_0(f"Skipped iterations {start} to {end} due to --skip-train-iteration-range flag.")
+            iteration_for_skipping = args.iteration
+            while iteration_for_skipping + 1 <= end:
+                try:
+                    _ = next(train_data_iterator)
+                except TypeError:
+                    pass
+                iteration_for_skipping += 1
+            continue
+
         update_num_microbatches(args.consumed_train_samples)
         if args.deepspeed:
             # inform deepspeed of any batch size changes

tests/ds_config_cl.json

@@ -20,8 +20,8 @@
     "max_difficulty": 128,
     "schedule_type": "fixed_linear",
     "schedule_config": {
-      "total_curriculum_step": 25,
-      "difficulty_step": 2
+      "total_curriculum_step": 30,
+      "difficulty_step": 4
     }
   },
   "steps_per_print": 2000,

tests/test_training.py

@@ -17,6 +17,7 @@
 import json
 import os
 import glob
+import re
 import unittest
 from pathlib import Path
 from parameterized import parameterized
@@ -78,7 +79,7 @@ def setUp(self):
         if os.path.exists(meg_lock_file_path):
             os.unlink(meg_lock_file_path)
 
-    def get_variation_config(self, variation, output_dir):
+    def get_variation_config(self, variation, output_dir, n_samples=None):
         data_dir = f"{self.data_dir}/gpt2"
 
         pp_size, tp_size, dp_size = get_3d_dimensions()
@@ -89,7 +90,9 @@ def get_variation_config(self, variation, output_dir):
             # we want to make sure at least tp=2 is used, so we swap tp and pp
             pp_size, tp_size = tp_size, pp_size
 
-        n_samples = 300 # about 56 iterations
+        if n_samples is None:
+            n_samples = 300 # about 56 iterations
+
         exit_interval = 20 # some samples in the first half and then some more in the 2nd half after resume
         seq_len = 128
 
@@ -107,7 +110,7 @@ def get_variation_config(self, variation, output_dir):
                 --pipeline-model-parallel-size {pp_size}
                 --distributed-backend nccl
 
-                --log-interval 5
+                --log-interval 1
                 --save-interval 10
                 --eval-interval 10
                 --eval-iters 5
@@ -219,6 +222,7 @@ def get_variation_config(self, variation, output_dir):
             new_ds_args = f"""
                 --deepspeed_config {self.test_file_dir_str}/ds_config.json
             """.split()
+
         else:
             raise ValueError(f"Don't know of variation {variation}")
 
@@ -495,3 +499,30 @@ def test_mode2_dataloading(self, variation):
         # test tensorboard
         tensorboard_files = glob.glob(f"{output_dir}/tensorboard/events*")
         self.assertEqual(len(tensorboard_files), 1, "tensorboard files")
+
+    def test_skip_train_iteration(self):
+        # skip iterations setup
+        extra_args = f"""
+            --skip-train-iteration-range 2-2 4-7
+        """.split()
+
+        src_dir = self.src_dir
+        output_dir = self.get_auto_remove_tmp_dir()
+        args, ds_args, num_gpus = self.get_variation_config("base", output_dir, n_samples=200)
+        args.extend(extra_args)
+        script = [f"{src_dir}/pretrain_gpt.py"]
+        launcher = get_launcher(num_gpus)
+        cmd = launcher + script + args + ds_args
+        # keep for quick debug
+        # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
+
+        with CaptureStdout() as cs:
+            execute_subprocess_async(cmd, env=self.get_env())
+
+        # check skipped iterations
+        self.assertIn("Skipped iterations 2 to 2 due to --skip-train-iteration-range flag", cs.out)
+        self.assertIn("Skipped iterations 4 to 7 due to --skip-train-iteration-range flag", cs.out)
+
+        train_iterations = range(1,10)
+        for i in train_iterations:
+            self.assertTrue(f"iteration {i:8d}/" in cs.out)