test: update llama tests

Author: cwher

tests/kit/model_zoo/transformers/llama.py

@@ -49,9 +49,9 @@ def data_gen_for_casual_lm():
     loss_fn_for_seq_classification = lambda output: output["logits"].mean()
 
     config = LlamaConfig(
-        num_hidden_layers=4,
-        hidden_size=128,
-        intermediate_size=256,
+        num_hidden_layers=8,
+        hidden_size=32,
+        intermediate_size=64,
         num_attention_heads=4,
         max_position_embeddings=128,
         num_labels=16,

tests/test_shardformer/test_model/test_shard_llama.py

@@ -5,6 +5,7 @@
 
 import colossalai
 from colossalai.logging import disable_existing_loggers
+from colossalai.shardformer import AdvancedPipelineConfig
 from colossalai.shardformer.layer.utils import Randomizer
 from colossalai.tensor.d_tensor.api import clear_layout_converter
 from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
@@ -24,9 +25,13 @@
 
 
 def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config):
+    enable_gradient_checkpointing = test_config.pop("enable_gradient_checkpointing", False)
     org_model, org_optimizer, sharded_model, sharded_optimizer, criterion, booster = build_model_from_hybrid_plugin(
         model_fn, loss_fn, test_config
     )
+    if enable_gradient_checkpointing:
+        org_model.gradient_checkpointing_enable()
+        sharded_model.unwrap().gradient_checkpointing_enable()
 
     org_loss, org_output, sharded_loss, sharded_output = run_forward_backward_with_hybrid_plugin(
         org_model, sharded_model, sharded_optimizer, data_gen_fn, output_transform_fn, criterion, booster
@@ -101,14 +106,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
             "use_lazy_init": True,
             "precision": "fp16",
             "initial_scale": 1,
-            "gradient_checkpointing_ratio": 0.5,
+            "enable_gradient_checkpointing": True,
+            "advanced_pipeline_config": AdvancedPipelineConfig(gradient_checkpointing_ratio=0.5),
         },
         {
             "tp_size": 1,
             "pp_size": 2,
             "num_microbatches": 4,
             "use_lazy_init": False,
             "precision": "fp32",
+            "enable_gradient_checkpointing": True,
+            "advanced_pipeline_config": AdvancedPipelineConfig(
+                num_stages=2, num_model_chunks=1, num_model_layers=8, num_layers_per_stage=[5, 3]
+            ),
         },
         {
             "tp_size": 4,
@@ -190,6 +200,14 @@ def run_llama_test(test_config):
             "precision": "fp16",
             "zero_stage": 1,
             "initial_scale": 1,
+            "enable_gradient_checkpointing": True,
+            "advanced_pipeline_config": AdvancedPipelineConfig(
+                num_stages=2,
+                num_model_chunks=2,
+                num_model_layers=8,
+                num_layers_per_stage=[3, 3, 1, 1],
+                num_ckpt_layers_per_stage=[0, 0, 1, 1],
+            ),
         },
     ],
 )