github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-28 21:17:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

[utils] refactor parallel layers checkpoint and bcast model on loading checkpoint (#1548)

Browse Source 
* refactor parallel layer

* broadcast rank0 model after load ckpt
This commit is contained in:
ver217
2022-09-06 20:18:35 +08:00
committed by GitHub
parent 2bed096848
commit ae71036cd2
6 changed files with 131 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
colossalai/nn/layer/base_layer.py
View File

@@ -5,9 +5,11 @@ import torch.nn as nn
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from contextlib import contextmanager
class ParallelLayer(nn.Module):
global_state_dict: bool = True
def __init__(self):
super().__init__()
@@ -26,10 +28,35 @@ class ParallelLayer(nn.Module):
self.pipeline_parallel_size = 1 if not gpc.is_initialized(ParallelMode.PIPELINE) else gpc.get_world_size(
ParallelMode.PIPELINE)
def _load_from_global_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
error_msgs):
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
error_msgs)
def _save_to_global_state_dict(self, destination, prefix, keep_vars):
return super()._save_to_state_dict(destination, prefix, keep_vars)
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
error_msgs):
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
error_msgs)
if gpc.get_local_rank(ParallelMode.TENSOR) != 0:
missing_keys.clear()
unexpected_keys.clear()
if self.global_state_dict:
if gpc.get_local_rank(ParallelMode.TENSOR) != 0:
missing_keys.clear()
unexpected_keys.clear()
return self._load_from_global_state_dict(state_dict, prefix, local_metadata, strict, missing_keys,
unexpected_keys, error_msgs)
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
error_msgs)
def _save_to_state_dict(self, destination, prefix, keep_vars):
if self.global_state_dict:
return self._save_to_global_state_dict(destination, prefix, keep_vars)
return super()._save_to_state_dict(destination, prefix, keep_vars)
@classmethod
@contextmanager
def use_local_state_dict(cls):
try:
cls.global_state_dict = False
yield
finally:
cls.global_state_dict = True