import os

import autoconfig.training_config as tc
import pytest
from omegaconf import OmegaConf


class TestCalculateTpPpMbsGrid:

    margin = 0.05

    @pytest.mark.parametrize(
        "model_size,layers,seq_length,model_name,train_cfg,expected",
        [
            # GPT-3 tests
            (
                0.126,
                12,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [2, 4, 8],
                    "micro_batch_sizes": [4, 8, 32],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": 1,
                    "max_model_parallel_size": 32,
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [2, 4, 8],
                    "mbs": [4, 8, 32],
                    "min_par": 1,
                    "max_par": 32,
                },
            ),
            (
                0.126,
                12,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.126,
                12,
                8192,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1, 2],
                    "mbs": [1, 2, 4],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.843,
                12,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.843,
                12,
                16384,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {"tp": [2, 4], "pp": [1], "mbs": [1, 2], "min_par": 1, "max_par": 8},
            ),
            (
                0.843,
                12,
                32768,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {"tp": [2, 4], "pp": [1, 2], "mbs": [1], "min_par": 1, "max_par": 8},
            ),
            (
                2,
                12,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                2,
                12,
                8192,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1, 2],
                    "mbs": [1, 2, 4],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                2.5,
                24,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                5.0,
                24,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                5.0,
                24,
                16384,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {"tp": [2, 4], "pp": [1, 2], "mbs": [1], "min_par": 1, "max_par": 8},
            ),
            (
                8.0,
                24,
                4096,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1, 2],
                    "mbs": [1, 2, 4],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                10.0,
                24,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4, 8],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                10.0,
                24,
                32768,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {"tp": [4, 8], "pp": [1, 2], "mbs": [1], "min_par": 4, "max_par": 16},
            ),
            (
                20.0,
                44,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1, 2, 4],
                    "mbs": [1, 2, 4],
                    "min_par": 4,
                    "max_par": 8,
                },
            ),
            (
                20.0,
                44,
                8192,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1, 2, 4],
                    "mbs": [1],
                    "min_par": 8,
                    "max_par": 32,
                },
            ),
            (
                43.0,
                52,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [2, 4, 8],
                    "pp": [1, 2, 4],
                    "mbs": [1, 2, 4],
                    "min_par": 8,
                    "max_par": 32,
                },
            ),
            (
                43.0,
                52,
                4096,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [2, 4],
                    "mbs": [1, 2],
                    "min_par": 8,
                    "max_par": 32,
                },
            ),
            (
                175.0,
                96,
                2048,
                "gpt3",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [8],
                    "pp": [4, 6, 8, 12, 16],
                    "mbs": [1, 2, 4],
                    "min_par": 32,
                    "max_par": 256,
                },
            ),
            # Llama tests
            (
                7,
                32,
                4096,
                "llama",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [1, 2, 4, 8],
                    "micro_batch_sizes": [1, 2, 4, 8],
                    "context_parallel_sizes": [1, 2],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": 1,
                    "max_model_parallel_size": 32,
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [1, 2, 4, 8],
                    "cp": [1, 2],
                    "mbs": [1, 2, 4, 8],
                    "min_par": 1,
                    "max_par": 32,
                },
            ),
            (
                8,
                32,
                8192,
                "llama",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [1, 2, 4, 8],
                    "micro_batch_sizes": [1, 2, 4, 8],
                    "context_parallel_sizes": [1, 2],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": 1,
                    "max_model_parallel_size": 40,
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [1, 2, 4, 8],
                    "cp": [1, 2],
                    "mbs": [1, 2, 4, 8],
                    "min_par": 1,
                    "max_par": 40,
                },
            ),
            # Mixtral tests
            (
                7,
                32,
                4096,
                "mixtral",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [1, 2, 4, 8],
                    "expert_parallel_sizes": [1, 2],
                    "micro_batch_sizes": [1, 2, 4, 8],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": 1,
                    "max_model_parallel_size": 32,
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [1, 2, 4, 8],
                    "ep": [1, 2],
                    "mbs": [1, 2, 4, 8],
                    "min_par": 1,
                    "max_par": 32,
                },
            ),
            # T5 tests
            (
                0.22,
                12,
                512,
                "t5",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [2, 4, 8],
                    "micro_batch_sizes": [4, 8, 32],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [2, 4, 8],
                    "mbs": [4, 8, 32],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.22,
                12,
                512,
                "t5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1],
                    "mbs": [16, 32, 64, 128],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                3.0,
                24,
                512,
                "t5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1],
                    "mbs": [4, 6, 8, 12, 16, 24, 32, 48],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                11.0,
                24,
                512,
                "t5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1],
                    "mbs": [2, 4, 6, 8, 12, 16, 24],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                23.0,
                36,
                512,
                "t5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1, 2],
                    "mbs": [1, 2, 4, 6, 8],
                    "min_par": 4,
                    "max_par": 16,
                },
            ),
            (
                41.0,
                48,
                512,
                "t5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1, 2, 4],
                    "mbs": [1, 2, 4, 6, 8],
                    "min_par": 8,
                    "max_par": 32,
                },
            ),
            # mT5 tests
            (
                0.17,
                6,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [2, 4, 8],
                    "micro_batch_sizes": [4, 8, 32],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [2, 4, 8],
                    "mbs": [4, 8, 32],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.17,
                6,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1],
                    "mbs": [16, 32, 64, 128],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.39,
                12,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1],
                    "mbs": [16, 32, 64, 128],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                3.2,
                24,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1],
                    "mbs": [4, 6, 8, 12, 16, 24, 32, 48],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                11.9,
                24,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1],
                    "mbs": [2, 4, 6, 8, 12, 16, 24],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                24.65,
                36,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1, 2],
                    "mbs": [1, 2, 4, 6, 8],
                    "min_par": 4,
                    "max_par": 16,
                },
            ),
            (
                42.54,
                48,
                512,
                "mt5",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1, 2, 4],
                    "mbs": [1, 2, 4, 6, 8],
                    "min_par": 8,
                    "max_par": 32,
                },
            ),
            # BERT tests
            (
                0.11,
                12,
                512,
                "bert",
                {
                    "tensor_parallel_sizes": [1, 2, 4, 5],
                    "pipeline_parallel_sizes": [2, 4, 8],
                    "micro_batch_sizes": [4, 8, 32],
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4, 5],
                    "pp": [2, 4, 8],
                    "mbs": [4, 8, 32],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                0.11,
                12,
                512,
                "bert",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                2.5,
                24,
                512,
                "bert",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [1, 2, 4],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 1,
                    "max_par": 8,
                },
            ),
            (
                5.0,
                24,
                512,
                "bert",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [2, 4, 8],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6, 8],
                    "min_par": 2,
                    "max_par": 8,
                },
            ),
            (
                10.0,
                24,
                512,
                "bert",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [2, 4, 8],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4, 6],
                    "min_par": 2,
                    "max_par": 8,
                },
            ),
            (
                20.0,
                48,
                512,
                "bert",
                {
                    "tensor_parallel_sizes": "auto",
                    "pipeline_parallel_sizes": "auto",
                    "micro_batch_sizes": "auto",
                    "gpu_memory_gb": 80,
                    "min_model_parallel_size": "auto",
                    "max_model_parallel_size": "auto",
                },
                {
                    "tp": [4, 8],
                    "pp": [1],
                    "mbs": [1, 2, 3, 4],
                    "min_par": 4,
                    "max_par": 8,
                },
            ),
        ],
    )
    def test_calculate_tp_pp_mbs_grid(
        self, model_size, layers, seq_length, model_name, train_cfg, expected
    ):
        params = {
            "model_size_in_b": model_size,
            "num_layers": layers,
            "seq_length": seq_length,
            "model_name": model_name,
            "train_cfg": train_cfg,
        }
        tp, pp, cp, ep, mbs, min_par, max_par = tc._calculate_tp_pp_mbs_grid(**params)

        assert tp == expected["tp"], f"TP should be {expected['tp']} but it is {tp}."
        assert pp == expected["pp"], f"PP should be {expected['pp']} but it is {pp}."

        if "cp" in expected:
            assert (
                cp == expected["cp"]
            ), f"PP should be {expected['cp']} but it is {cp}."
        if "ep" in expected:
            assert (
                ep == expected["ep"]
            ), f"PP should be {expected['ep']} but it is {ep}."

        assert (
            mbs == expected["mbs"]
        ), f"MBS should be {expected['mbs']} but it is {mbs}."
        assert (
            min_par == expected["min_par"]
        ), f"Minimum paralellism should be {expected['min_par']} but it is {min_par}."
        assert (
            max_par == expected["max_par"]
        ), f"Minimum paralellism should be {expected['max_par']} but it is {max_par}."
