Merge pull request #6505 from hpcaitech/feat/numerics-model-alignment

[test] numerical alignment: whole models through booster plugins, data-parallel and resume equivalence, acceptance report

Author: yuxuan-lou

.github/scripts/ci/run_colossalai_batch.sh

@@ -76,7 +76,11 @@ case "${BATCH}" in
             tests/test_config/test_load_config.py
             tests/test_infer/test_async_engine/test_request_tracer.py
             tests/test_lazy/test_chatglm_cpu.py
+            tests/test_numerics/test_numerics_acceptance.py
             tests/test_numerics/test_numerics_core.py
+            tests/test_numerics/test_numerics_data_parallel.py
+            tests/test_numerics/test_numerics_independent_impl.py
+            tests/test_numerics/test_numerics_resume.py
             tests/test_optimizer/test_lr_scheduler.py
             tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py
             tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py
@@ -166,7 +170,7 @@ case "${BATCH}" in
         ;;
     5|05)
         LABEL="05-booster-lora-2gpu"
-        DESCRIPTION="DDP, FSDP, LowLevelZeRO and LoRA"
+        DESCRIPTION="DDP, FSDP, LowLevelZeRO, LoRA and the numerical alignment tool on its plugins"
         EXPECTED_GPUS=2
         DEFAULT_TIMEOUT_MIN=90
         NEEDS_CUDA_TOOLKIT=1
@@ -176,6 +180,7 @@ case "${BATCH}" in
             tests/test_booster/test_plugin/test_torch_fsdp_plugin.py
             tests/test_booster/test_plugin/test_low_level_zero_plugin.py
             tests/test_lora/test_lora.py
+            tests/test_numerics
         )
         ;;
     6|06)

.github/scripts/ci/run_cpu_unit_tests.sh

@@ -14,7 +14,11 @@ TESTS=(
     tests/test_config/test_load_config.py
     tests/test_infer/test_async_engine/test_request_tracer.py
     tests/test_lazy/test_chatglm_cpu.py
+    tests/test_numerics/test_numerics_acceptance.py
     tests/test_numerics/test_numerics_core.py
+    tests/test_numerics/test_numerics_data_parallel.py
+    tests/test_numerics/test_numerics_independent_impl.py
+    tests/test_numerics/test_numerics_resume.py
     tests/test_optimizer/test_lr_scheduler.py
     tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py
     tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py

tests/kit/model_zoo/registry.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python
 from dataclasses import dataclass
-from typing import Callable, List, Union
+from typing import Callable, List, Optional, Union
+
+from tests.kit.numerics.spec import NumericsSpec
 
 __all__ = ["ModelZooRegistry", "ModelAttribute", "model_zoo"]
 
@@ -13,10 +15,13 @@ class ModelAttribute:
     Args:
         has_control_flow (bool): Whether the model contains branching in its forward method.
         has_stochastic_depth_prob (bool): Whether the model contains stochastic depth probability. Often seen in the torchvision models.
+        numerics (NumericsSpec): How to check the model against its reference implementation. Models without one are
+            not part of whole-model alignment. See ``tests/kit/numerics/spec.py``.
     """
 
     has_control_flow: bool = False
     has_stochastic_depth_prob: bool = False
+    numerics: Optional[NumericsSpec] = None
 
 
 class ModelZooRegistry(dict):

tests/kit/model_zoo/transformers/mixtral.py

@@ -3,6 +3,8 @@
 import transformers
 from transformers import MixtralConfig
 
+from tests.kit.numerics.spec import NumericsSpec, Routers
+
 from ..registry import ModelAttribute, model_zoo
 
 # ===============================
@@ -57,20 +59,31 @@ def data_gen_for_sequence_classification():
     # FlashAttention2.  ShardFormer tests enable their optimized attention
     # paths explicitly, so keep this baseline self-contained.
     attn_implementation="eager",
+    # The default grouped_mm experts kernel has no float64 path, which the
+    # float64 ground truth of numerical alignment (tests/kit/numerics) needs.
+    experts_implementation="eager",
     torch_dtype="float16",
     output_router_logits=True,
 )
 
 if hasattr(config, "pad_token_id"):
     config.pad_token_id = config.eos_token_id
 
+
+def is_router(name, module):
+    # the gate of each sparse MoE block; its first output is the router logits
+    return name.endswith("mlp.gate")
+
+
+numerics = NumericsSpec(routers=Routers(select=is_router, top_k=config.num_experts_per_tok))
+
 model_zoo.register(
     name="transformers_mixtral",
     model_fn=lambda: transformers.MixtralModel(config),
     data_gen_fn=data_gen,
     output_transform_fn=output_transform_fn,
     loss_fn=loss_fn_for_mixtral_model,
-    model_attribute=ModelAttribute(has_control_flow=True),
+    model_attribute=ModelAttribute(has_control_flow=True, numerics=numerics),
 )
 # model_zoo.register(
 #     name="transformers_mixtral_for_casual_lm",

tests/kit/numerics/README.md

@@ -1,19 +1,55 @@
-from .capture import cast_floating, full_precision_matmul, run_fwd_bwd
+from .capture import (
+    RoutingRecorder,
+    cast_floating,
+    count_calls,
+    full_precision_matmul,
+    missing_grads,
+    run_fwd_bwd,
+    run_optimizer,
+    upstream_loss,
+)
 from .compare import TensorMetrics, flatten_tensors, tensor_metrics
-from .judge import AlignmentReport, Entry, assert_aligned, error_multiplier, judge
+from .judge import (
+    AlignmentReport,
+    Entry,
+    assert_aligned,
+    diagnose_routing,
+    error_multiplier,
+    judge,
+    judge_equal,
+    judge_relative_bound,
+    machine_epsilon,
+    relative_bound,
+    require_gradients,
+)
+from .spec import NumericsSpec, Routers, TokenLoss
 from .trace import Trace
 
 __all__ = [
     "AlignmentReport",
     "Entry",
+    "NumericsSpec",
+    "Routers",
+    "RoutingRecorder",
     "TensorMetrics",
+    "TokenLoss",
     "Trace",
     "assert_aligned",
     "cast_floating",
+    "count_calls",
+    "diagnose_routing",
     "error_multiplier",
     "flatten_tensors",
     "full_precision_matmul",
     "judge",
+    "judge_equal",
+    "judge_relative_bound",
+    "machine_epsilon",
+    "missing_grads",
+    "relative_bound",
+    "require_gradients",
     "run_fwd_bwd",
+    "run_optimizer",
     "tensor_metrics",
+    "upstream_loss",
 ]