test_megatron_gpt.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

"""Tests for BERT model benchmarks."""

import os
from pathlib import Path
import shlex
import statistics
from unittest import mock
import unittest
from superbench.benchmarks.context import ModelAction, Precision

from tests.helper import decorator
from superbench.benchmarks import BenchmarkRegistry, Platform, ReturnCode
from tests.helper.testcase import BenchmarkTestCase


def normalize_command(cmd):
    """Convert a CLI string into a list of meaningful argument units (key-value or flag)."""
    tokens = shlex.split(cmd)
    units = []
    i = 0
    while i < len(tokens):
        if tokens[i].startswith('--'):
            if i + 1 >= len(tokens) or tokens[i + 1].startswith('--'):
                units.append(tokens[i])    # flag-only
                i += 1
            else:
                units.append(f'{tokens[i]} {tokens[i + 1]}')    # key-value pair
                i += 2
        else:
            # Include positional args like torchrun, script path, etc.
            units.append(tokens[i])
            i += 1
    return sorted(units)


class MegatronGPTTest(BenchmarkTestCase, unittest.TestCase):
    """Tests for IBBenchmark benchmark."""
    @classmethod
    def setUpClass(cls):
        """Hook method for setting up class fixture before running tests in the class."""
        super().setUpClass()
        cls.benchmark_name = 'megatron-gpt'
        cls.createMockEnvs(cls)
        cls.hostfile_path = os.path.join(cls._tmp_dir, 'hostfile')

    @classmethod
    def tearDownClass(cls):
        """Hook method for deconstructing the class fixture after running all tests in the class."""
        for p in [
            Path(cls._tmp_dir) / 'pretrain_gpt.py',
            Path(cls._tmp_dir) / 'customdataset_text_document.bin',
            Path(cls._tmp_dir) / 'customdataset_text_document.idx',
            Path(cls._tmp_dir) / 'hostfile'
        ]:
            if p.is_file():
                p.unlink()
        super().tearDownClass()

    @mock.patch('superbench.benchmarks.model_benchmarks.MegatronGPT._generate_dataset')
    def test_megatron_gpt_preprocess(self, mock_generate_dataset):
        """Test megatron-gpt benchmark."""
        # Check registry.
        (benchmark_cls, _) = BenchmarkRegistry._BenchmarkRegistry__select_benchmark(self.benchmark_name, Platform.CUDA)
        assert (benchmark_cls)
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--hostfile {self.hostfile_path} --batch_size 2048',
        )

        # Check init distribued setting.
        os.environ['OMPI_COMM_WORLD_SIZE'] = '2'
        os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_RANK'] = '0'
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = '12345'
        with open(self.hostfile_path, 'w') as f:
            f.write('host1\n')
            f.write('host2\n')
            f.write('host3\n')
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is False)
        assert (benchmark.return_code == ReturnCode.DISTRIBUTED_SETTING_INIT_FAILURE)

        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters='--hostfile xxx --batch_size 2048',
        )
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is False)
        assert (benchmark.return_code == ReturnCode.DISTRIBUTED_SETTING_INIT_FAILURE)

        os.environ['OMPI_COMM_WORLD_SIZE'] = '3'
        os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'] = '1'
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--hostfile {self.hostfile_path} --batch_size 2048',
        )
        mock_generate_dataset.return_value = True
        benchmark._preprocess()
        self.assertEqual(benchmark._num_nodes, 3)
        self.assertEqual(
            benchmark._distributed_args,
            '--nproc_per_node {0} --nnodes {1} --node_rank {2} --master_addr {3} --master_port {4}'.format(
                benchmark._args.num_gpus, benchmark._num_nodes, 0, 'localhost', '12345'
            )
        )

        # Check preprocessing.
        # Negative cases
        # no code_base
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} --batch_size 2048',
        )
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is False)
        self.createMockFiles(['pretrain_gpt.py'])
        # invalid micro batch size
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} --micro_batch_size -1',
        )
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is False)
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} --micro_batch_size 4096',
        )
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is False)
        # invalid precision
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} \
                --batch_size 2048 --precision int8',
        )
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is False)
        # Positive cases
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} --batch_size 2048',
        )
        mock_generate_dataset.return_value = True
        ret = benchmark._preprocess()
        assert (ret is True)

    def test_megatron_gpt_dataset(self):
        """Test dataset genreation."""
        (benchmark_cls, _) = BenchmarkRegistry._BenchmarkRegistry__select_benchmark(self.benchmark_name, Platform.CUDA)
        assert (benchmark_cls)
        os.environ['OMPI_COMM_WORLD_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_RANK'] = '0'
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = '12345'
        # use existing dataset
        self.createMockFiles(['customdataset_text_document.bin', 'customdataset_text_document.idx'])
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base /root/Megatron-DeepSpeed --data_home {self._tmp_dir} \
                --batch_size 2048 --data_prefix customdataset_text_document',
        )
        ret = benchmark._preprocess()
        ret = benchmark._generate_dataset()
        assert (ret is True)

    @mock.patch('superbench.benchmarks.model_benchmarks.MegatronGPT._generate_dataset')
    def test_megatron_gpt_command(self, mock_generate_dataset):
        """Test command generation."""
        (benchmark_cls, _) = BenchmarkRegistry._BenchmarkRegistry__select_benchmark(self.benchmark_name, Platform.CUDA)
        assert (benchmark_cls)
        os.environ['OMPI_COMM_WORLD_SIZE'] = '2'
        os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_RANK'] = '0'
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = '12345'
        with open(self.hostfile_path, 'w') as f:
            f.write('host1\n')
            f.write('host2\n')
        # use url to process dataset
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} \
                --num_warmup 0 --num_steps 10 --batch_size 2048 --data_prefix dataset_text_document \
                --override_opt_param_scheduler',
        )
        mock_generate_dataset.return_value = True
        benchmark._preprocess()
        benchmark._data_options = f'\
            --vocab-file {self._tmp_dir}/gpt2-vocab.json \
            --merge-file {self._tmp_dir}/gpt2-merges.txt \
            --data-path {self._tmp_dir}/dataset_text_document \
            --split 949,50,1'

        script_path = str(Path(self._tmp_dir) / 'pretrain_gpt.py')
        expected_command_template = 'torchrun {distributed_args} {script_path} \
            --tokenizer-type GPT2BPETokenizer \
            --override-opt_param-scheduler \
            --adam-beta1 0.9 \
            --adam-beta2 0.95 \
            --tensor-model-parallel-size 1 \
            --init-method-std 0.009 \
            --lr-decay-samples 43945312 \
            --lr-warmup-samples 0 \
            --lr-decay-style cosine \
            --micro-batch-size 2 \
            --global-batch-size 2048 \
            --num-layers 32 \
            --hidden-size 4096 \
            --num-attention-heads 32 \
            --seq-length 2048 \
            --max-position-embeddings 2048 \
            --train-samples 20480 \
            --lr 0.00012 \
            --min-lr 1e-06 \
            --log-interval 1 \
            --eval-interval 10 \
            --eval-iters 0 \
            --save-interval 10000 \
            --weight-decay 0.1 \
            --clip-grad 1.0 \
            --hysteresis 2 \
            --num-workers 8 \
            --attention-dropout 0.0 \
            --hidden-dropout 0.0 \
            --optimizer adam \
            --use-distributed-optimizer \
            {precision} \
            --seed 1234 \
            --log-throughput {data_options}'

        precision = Precision.FLOAT32
        expected_command = expected_command_template.format(
            precision='',
            data_options=benchmark._data_options,
            distributed_args=benchmark._distributed_args,
            script_path=script_path
        )
        command = benchmark._megatron_command(precision)
        actual_units = normalize_command(command)
        expected_units = normalize_command(expected_command)
        self.assertEqual(actual_units, expected_units)

        precision = Precision.FLOAT16
        expected_command = expected_command_template.format(
            precision='--fp16',
            data_options=benchmark._data_options,
            distributed_args=benchmark._distributed_args,
            script_path=script_path
        )
        command = benchmark._megatron_command(precision)
        actual_units = normalize_command(command)
        expected_units = normalize_command(expected_command)
        self.assertEqual(actual_units, expected_units)

        precision = Precision.BFLOAT16
        expected_command = expected_command_template.format(
            precision='--bf16',
            data_options=benchmark._data_options,
            distributed_args=benchmark._distributed_args,
            script_path=script_path
        )
        command = benchmark._megatron_command(precision)
        actual_units = normalize_command(command)
        expected_units = normalize_command(expected_command)
        self.assertEqual(actual_units, expected_units)

        os.environ['OMPI_COMM_WORLD_SIZE'] = '1'
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} \
                --num_warmup 0 --num_steps 10 --batch_size 2048 --data_prefix dataset_text_document \
                --deepspeed --override_opt_param_scheduler',
        )
        benchmark._preprocess()
        benchmark._data_options = f'\
            --vocab-file {self._tmp_dir}/gpt2-vocab.json \
            --merge-file {self._tmp_dir}/gpt2-merges.txt \
            --data-path {self._tmp_dir}/dataset_text_document \
            --split 949,50,1'

        command = benchmark._megatron_command(Precision.BFLOAT16)
        expected_command = 'deepspeed {script_path} --override-opt_param-scheduler \
            --tokenizer-type GPT2BPETokenizer \
            --adam-beta1 0.9 \
            --adam-beta2 0.95 \
            --tensor-model-parallel-size 1 \
            --init-method-std 0.009 \
            --lr-decay-samples 43945312 \
            --lr-warmup-samples 0 \
            --lr-decay-style cosine \
            --micro-batch-size 2 \
            --global-batch-size 2048 \
            --num-layers 32 \
            --hidden-size 4096 \
            --num-attention-heads 32 \
            --seq-length 2048 \
            --max-position-embeddings 2048 \
            --train-samples 20480 \
            --lr 0.00012 \
            --min-lr 1e-06 \
            --log-interval 1 \
            --eval-interval 10 \
            --eval-iters 0 \
            --save-interval 10000 \
            --weight-decay 0.1 \
            --clip-grad 1.0 \
            --hysteresis 2 \
            --num-workers 8 \
            --attention-dropout 0.0 \
            --hidden-dropout 0.0 \
            --optimizer adam \
            --use-distributed-optimizer \
            {precision} \
            --seed 1234 {data_options} {deepseed_options}'

        expect_ds_options = f'\
            --deepspeed \
            --deepspeed_config {benchmark._config_json_path} \
            --zero-stage 1 \
            --pipeline-model-parallel-size 1 \
            --train-tokens 300000000000 \
            --data-impl mmap --no-pipeline-parallel'

        expected_command = expected_command.format(
            precision='--bf16',
            data_options=benchmark._data_options,
            deepseed_options=expect_ds_options,
            script_path=script_path
        )
        command = benchmark._megatron_command(Precision.BFLOAT16)
        actual_units = normalize_command(command)
        expected_units = normalize_command(expected_command)
        self.assertEqual(actual_units, expected_units)

    def test_deepseek_v2_command(self):
        """Test v2 command."""
        # test deepspeed with megatron
        os.environ['OMPI_COMM_WORLD_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_RANK'] = '0'
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = '12345'
        with open(self.hostfile_path, 'w') as f:
            f.write('host1\n')

        benchmark_name = 'megatron-deepseek-v2'
        (benchmark_cls, _) = BenchmarkRegistry._BenchmarkRegistry__select_benchmark(self.benchmark_name, Platform.ROCM)
        assert (benchmark_cls)
        benchmark = benchmark_cls(
            benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --hostfile {self.hostfile_path} '
            '--num_warmup 0 '
            '--num_steps 10 '
            '--batch_size 256 '
            '--expert_model_parallel_size 8 '
            '--micro_batch_size 2 '
            '--mock_data '
            '--model=deepseek '
            '--tokenizer_type=DeepSeekV2Tokenizer '
            '--transformer_impl=transformer_engine '
            '--num_layers=27 '
            '--hidden_size=1024 '
            '--seq_len=4096 '
            '--ffn_hidden_size=10944 '
            '--num_attn_heads=16 '
            '--moe_ffn_hidden_size=1408 '
            '--enable_shared_expert '
            '--moe_layer_freq=1 '
            '--num_shared_experts=2 '
            '--moe_router_topk=6 '
            '--moe_aux_loss_coeff=0.01 '
            '--moe_router_load_balancing_type=aux_loss '
            '--num_experts=64 '
            '--patch_tokenizer_type=DeepSeekV2Tokenizer '
            '--position_embedding_type=rope '
            '--no_rope_fusion '
            '--rotary_base=10000 '
            '--rotary_scaling_factor=40 '
            '--qk_nope_head_dim=128 '
            '--qk_rope_head_dim=64 '
            '--v_head_dim=128 '
            '--ffn_hidden_size=10944 '
            '--swiglu '
            '--normalization=RMSNorm '
            '--norm_epsilon=1e-06 '
            '--no_bias_swiglu_fusion '
            '--disable_bias_linear '
            '--untie_embeddings_and_output_weights '
            '--extra_vocab_size=2400 '
            '--load=deepseek-ai/DeepSeek-V2-Lite '
            '--no_load_optim '
            '--no_load_rng '
            '--ckpt_format=torch '
            '--eod_mask_loss '
            '--train_mode=pretrain '
            '--data_cache_path=/root/cache '
            '--max_padding_length=4096 '
            '--kv_lora_rank=512 '
            '--dataloader_type=cyclic '
        )

        benchmark._preprocess()
        benchmark._data_options = '\
            --mock-data \
            --dataloader-type cyclic \
            --data-cache-path /root/cache \
            --dataset LLama-Pretrain-Idxmap'

        precision = Precision.BFLOAT16
        command = benchmark._megatron_command(precision)

        expected_command = (
            'torchrun {script_path} --bf16 \
            --init-method-std 0.009 \
            --adam-beta1 0.9 \
            --hidden-dropout 0.0 \
            --min-lr 1e-06 \
            --lr 0.00012 \
            --optimizer adam \
            --log-interval 1 \
            --eval-interval 10 \
            --seed 1234 \
            --eval-iters 0 \
            --max-position-embeddings 4096 \
            --hysteresis 2 \
            --lr-decay-style cosine \
            --lr-decay-samples 43945312 \
            --clip-grad 1.0 \
            --save-interval 10000 \
            --adam-beta2 0.95 \
            --moe-aux-loss-coeff 0.01 \
            --log-throughput \
            --num-workers 8 \
            --use-distributed-optimizer \
            --attention-dropout 0.0 \
            --tensor-model-parallel-size 1 \
            --lr-warmup-samples 0 \
            --weight-decay 0.1 \
            --train-samples 2560 \
            --no-load-optim \
            --load deepseek-ai/DeepSeek-V2-Lite \
            --no-load-rng \
            --ffn-hidden-size 10944 \
            --patch-tokenizer-type DeepSeekV2Tokenizer \
            --swiglu \
            --normalization RMSNorm \
            --norm-epsilon 1e-06 \
            --no-bias-swiglu-fusion \
            --no-rope-fusion \
            --position-embedding-type rope \
            --untie-embeddings-and-output-weights \
            --disable-bias-linear \
            --ckpt-format torch \
            --rotary-base 10000 \
            --rotary-scaling-factor 40 \
            --eod-mask-loss \
            --moe-ffn-hidden-size 1408 \
            --enable-shared-expert \
            --moe-layer-freq 1 \
            --num-shared-experts 2 \
            --moe-router-topk 6 \
            --kv-lora-rank 512 \
            --qk-nope-head-dim 128 \
            --qk-rope-head-dim 64 \
            --v-head-dim 128 \
            --moe-router-load-balancing-type aux_loss \
            --train-mode pretrain \
            --extra-vocab-size 2400 \
            --global-batch-size 256 \
            --micro-batch-size 2 \
            --num-layers 27 \
            --hidden-size 1024 \
            --seq-length 4096 \
            --num-attention-heads 16 \
            --tokenizer-type DeepSeekV2Tokenizer \
            --transformer-impl transformer_engine \
            --num-experts 64 \
            --expert-model-parallel-size 8 \
            --max-padding-length 4096 \
            {data_options} \
            {disitributed_args}'
        ).format(
            script_path=str(Path(self._tmp_dir) / 'pretrain_deepseek.py'),
            data_options=benchmark._data_options,
            disitributed_args=benchmark._distributed_args
        )
        actual_units = normalize_command(command)
        expected_units = normalize_command(expected_command)

        self.assertEqual(actual_units, expected_units)

    @decorator.load_data('tests/data/megatron_deepspeed.log')
    @mock.patch('superbench.benchmarks.model_benchmarks.MegatronGPT._generate_dataset')
    def test_megatron_parse_log(self, raw_output, mock_generate_dataset):
        """Test parse log function."""
        (benchmark_cls, _) = BenchmarkRegistry._BenchmarkRegistry__select_benchmark(self.benchmark_name, Platform.CUDA)
        assert (benchmark_cls)
        os.environ['OMPI_COMM_WORLD_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_LOCAL_SIZE'] = '1'
        os.environ['OMPI_COMM_WORLD_RANK'] = '0'
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = '12345'

        # use url to process dataset
        benchmark = benchmark_cls(
            self.benchmark_name,
            parameters=f'--code_base {self._tmp_dir} --num_warmup 0 --num_steps 10 --batch_size 2048',
        )
        mock_generate_dataset.return_value = True
        benchmark._preprocess()
        benchmark._data_options = f'\
            --vocab-file {self._tmp_dir}/gpt2-vocab.json \
            --merge-file {self._tmp_dir}/gpt2-merges.txt \
            --data-path {self._tmp_dir}/dataset_text_document \
            --data-impl mmap'

        iteration_times, tflops, mem_allocated, max_mem_allocated = benchmark._parse_log(raw_output)
        assert (statistics.mean(iteration_times) == 75239.24)
        assert (statistics.mean(tflops) == 149.136)
        assert (statistics.mean(mem_allocated) == 17.535637855529785)
        assert (statistics.mean(max_mem_allocated) == 66.9744234085083)

        info = {'tflops': tflops, 'mem_allocated': mem_allocated, 'max_mem_allocated': max_mem_allocated}
        benchmark._process_info(ModelAction.TRAIN, Precision.FLOAT16, info)
        assert (benchmark.result is not None)
        assert (benchmark.result['fp16_train_tflops'][0] == 149.136)
        assert (benchmark.result['fp16_train_mem_allocated'][0] == 17.535637855529785)
        assert (benchmark.result['fp16_train_max_mem_allocated'][0] == 66.9744234085083)