[booster] add tests for ddp and low level zero's checkpointio (#3715)

* [booster] update tests for booster * [booster] update tests for booster * [booster] update tests for booster * [booster] update tests for booster * [booster] update tests for booster * [booster] update booster tutorials#3717, fix recursive check
2025-09-08 20:40:34 +00:00 · 2023-05-10 12:17:02 +08:00
parent 6552cbf8e1
commit 20068ba188
6 changed files with 261 additions and 125 deletions
--- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
+++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
@@ -0,0 +1,98 @@
+import tempfile
+
+import pytest
+import torch
+
+import colossalai
+from colossalai.booster.plugin.gemini_plugin import GeminiCheckpointIO
+from colossalai.testing import check_state_dict_equal, parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.utils.cuda import get_current_device
+from colossalai.zero import ColoInitContext, ZeroDDP
+from colossalai.zero.gemini.chunk import ChunkManager, search_chunk_configuration
+from colossalai.zero.gemini.gemini_mgr import GeminiManager
+from tests.components_to_test.registry import non_distributed_component_funcs
+
+
+@parameterize('placement_policy', ['cuda', 'cpu'])
+@parameterize('model_name', ['bert'])
+@parameterize('use_safetensors', [True, False])
+def exam_state_dict_with_origin(placement_policy, model_name, use_safetensors: bool):
+    from transformers import BertForSequenceClassification
+
+    model_ckpt_dir = tempfile.TemporaryDirectory()
+    get_components_func = non_distributed_component_funcs.get_callable(model_name)
+    model_builder, *_ = get_components_func()
+    with ColoInitContext(device=(get_current_device())):
+        bert_model = model_builder()
+    bert_model.config.save_pretrained(save_directory=(model_ckpt_dir.name))
+
+    config_dict, *_ = search_chunk_configuration(bert_model, search_range_mb=1, search_interval_byte=100)
+    chunk_manager = ChunkManager(config_dict)
+    gemini_manager = GeminiManager(placement_policy, chunk_manager)
+    bert_model = ZeroDDP(bert_model, gemini_manager)
+    bert_model.train()
+
+    ckpt_io = GeminiCheckpointIO()
+    if ckpt_io.coordinator.is_master():
+        model_size = sum(p.numel() * p.element_size() for p in bert_model.parameters()) / 1024**2
+        ckpt_io.save_model(bert_model, (model_ckpt_dir.name),
+                           True,
+                           True,
+                           '', (model_size / 3),
+                           use_safetensors=use_safetensors)
+        new_bert_model = BertForSequenceClassification.from_pretrained(model_ckpt_dir.name)
+        check_state_dict_equal(bert_model.state_dict(only_rank_0=True, dtype=(torch.float32)),
+                               new_bert_model.state_dict(), False)
+    model_ckpt_dir.cleanup()
+
+
+@parameterize('placement_policy', ['cuda', 'cpu'])
+@parameterize('model_name', ['gpt2', 'bert'])
+@parameterize('use_safetensors', [True, False])
+def exam_state_dict(placement_policy, model_name: str, use_safetensors: bool):
+    get_components_func = non_distributed_component_funcs.get_callable(model_name)
+    model_builder, *_ = get_components_func()
+    with ColoInitContext(device=(get_current_device())):
+        model = model_builder()
+        new_model = model_builder()
+    config_dict, *_ = search_chunk_configuration(model, search_range_mb=1, search_interval_byte=100)
+    chunk_manager = ChunkManager(config_dict)
+    gemini_manager = GeminiManager(placement_policy, chunk_manager)
+    model = ZeroDDP(model, gemini_manager)
+
+    model.train()
+    #new model
+    new_config_dict, *_ = search_chunk_configuration(new_model, search_range_mb=1, search_interval_byte=100)
+    new_chunk_manager = ChunkManager(new_config_dict)
+    new_gemini_manager = GeminiManager(placement_policy, new_chunk_manager)
+    new_model = ZeroDDP(new_model, new_gemini_manager)
+
+    model_ckpt_dir = tempfile.TemporaryDirectory()
+    ckpt_io = GeminiCheckpointIO()
+    model_size = sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**2
+    ckpt_io.save_model(model, (model_ckpt_dir.name),
+                       True,
+                       True,
+                       'epoch', (model_size / 3),
+                       use_safetensors=use_safetensors)
+
+    if ckpt_io.coordinator.is_master():
+        ckpt_io.load_model(new_model, (model_ckpt_dir.name), strict=True)
+        model_dict = model.state_dict(only_rank_0=True)
+        new_model_dict = new_model.state_dict(only_rank_0=True)
+        check_state_dict_equal(model_dict, new_model_dict, False)
+    model_ckpt_dir.cleanup()
+
+
+def run_dist(rank, world_size, port):
+    config = {}
+    colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    exam_state_dict()
+    exam_state_dict_with_origin()
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [4, 4])
+@rerun_if_address_is_in_use()
+def test_gemini_ckpIO(world_size):
+    spawn(run_dist, world_size)