From e429d5e8230631a00a9777204348e3f59bdf5ffb Mon Sep 17 00:00:00 2001 From: Wenhao Chen Date: Tue, 28 Nov 2023 10:45:42 +0800 Subject: [PATCH 1/2] fix: remove unused code --- colossalai/shardformer/policies/gptj.py | 1 - 1 file changed, 1 deletion(-) diff --git a/colossalai/shardformer/policies/gptj.py b/colossalai/shardformer/policies/gptj.py index fe52b00fbebe..9feb826c4624 100644 --- a/colossalai/shardformer/policies/gptj.py +++ b/colossalai/shardformer/policies/gptj.py @@ -163,7 +163,6 @@ def get_held_layers(self) -> List[nn.Module]: layers_per_stage = self.distribute_layers(len(module.h), stage_manager.num_stages) if stage_manager.is_first_stage(): held_layers.append(module.wte) - # held_layers.append(module.wpe) held_layers.append(module.drop) start_idx, end_idx = self.get_stage_index(layers_per_stage, stage_manager.stage) held_layers.extend(module.h[start_idx:end_idx]) From 81fbec222738a503e27a16821bef388050c4b574 Mon Sep 17 00:00:00 2001 From: Wenhao Chen Date: Tue, 28 Nov 2023 11:23:05 +0800 Subject: [PATCH 2/2] feat: move colo_init_context to legacy --- colossalai/legacy/zero/gemini/__init__.py | 3 +++ .../{ => legacy}/zero/gemini/colo_init_context.py | 0 colossalai/zero/__init__.py | 11 +---------- colossalai/zero/gemini/__init__.py | 3 --- 4 files changed, 4 insertions(+), 13 deletions(-) rename colossalai/{ => legacy}/zero/gemini/colo_init_context.py (100%) diff --git a/colossalai/legacy/zero/gemini/__init__.py b/colossalai/legacy/zero/gemini/__init__.py index b272980d34d8..f30bccea4a95 100644 --- a/colossalai/legacy/zero/gemini/__init__.py +++ b/colossalai/legacy/zero/gemini/__init__.py @@ -1,3 +1,4 @@ +from .colo_init_context import ColoInitContext, post_process_colo_init_ctx from .ophooks import BaseOpHook, register_ophooks_recursively from .stateful_tensor import StatefulTensor from .stateful_tensor_mgr import StatefulTensorMgr @@ -11,4 +12,6 @@ "AutoTensorPlacementPolicy", "register_ophooks_recursively", "BaseOpHook", + "ColoInitContext", + "post_process_colo_init_ctx", ] diff --git a/colossalai/zero/gemini/colo_init_context.py b/colossalai/legacy/zero/gemini/colo_init_context.py similarity index 100% rename from colossalai/zero/gemini/colo_init_context.py rename to colossalai/legacy/zero/gemini/colo_init_context.py diff --git a/colossalai/zero/__init__.py b/colossalai/zero/__init__.py index 90d0f8de1916..5ad59e83234b 100644 --- a/colossalai/zero/__init__.py +++ b/colossalai/zero/__init__.py @@ -1,11 +1,4 @@ -from .gemini import ( - ColoInitContext, - GeminiAdamOptimizer, - GeminiDDP, - GeminiOptimizer, - get_static_torch_model, - post_process_colo_init_ctx, -) +from .gemini import GeminiAdamOptimizer, GeminiDDP, GeminiOptimizer, get_static_torch_model from .low_level import LowLevelZeroOptimizer from .wrapper import zero_model_wrapper, zero_optim_wrapper @@ -16,7 +9,5 @@ "zero_model_wrapper", "zero_optim_wrapper", "LowLevelZeroOptimizer", - "ColoInitContext", - "post_process_colo_init_ctx", "get_static_torch_model", ] diff --git a/colossalai/zero/gemini/__init__.py b/colossalai/zero/gemini/__init__.py index 358d5c7fd289..6d93ca8edfda 100644 --- a/colossalai/zero/gemini/__init__.py +++ b/colossalai/zero/gemini/__init__.py @@ -1,5 +1,4 @@ from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration -from .colo_init_context import ColoInitContext, post_process_colo_init_ctx from .gemini_ddp import GeminiDDP from .gemini_mgr import GeminiManager from .gemini_optimizer import GeminiAdamOptimizer, GeminiOptimizer @@ -15,6 +14,4 @@ "get_static_torch_model", "GeminiAdamOptimizer", "GeminiOptimizer", - "ColoInitContext", - "post_process_colo_init_ctx", ]