github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into feature/shardformer

Browse Source
This commit is contained in:
Hongxin Liu
2023-09-05 21:54:08 +08:00
committed by GitHub
parent bd18678478 ac178ca5c1
commit fae6c92ead
113 changed files with 629 additions and 633 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/zero/legacy/gemini/ophooks/_shard_grad_ophook.py
View File

@@ -1,6 +1,6 @@
import torch
from colossalai.registry import OPHOOKS
from colossalai.legacy.registry import OPHOOKS
from . import BaseOpHook

2
colossalai/zero/legacy/gemini/ophooks/_shard_param_ophook.py
View File

@@ -1,6 +1,6 @@
import torch
from colossalai.registry import OPHOOKS
from colossalai.legacy.registry import OPHOOKS
from . import BaseOpHook

2
colossalai/zero/legacy/sharded_model/zero_hook.py
View File

@@ -3,8 +3,8 @@ from typing import Optional
import torch
import torch.distributed as dist
from colossalai.legacy.registry import OPHOOKS
from colossalai.logging import get_dist_logger
from colossalai.registry import OPHOOKS
from colossalai.utils import get_current_device
from colossalai.zero.gemini.memory_tracer import MemStatsCollector
from colossalai.zero.legacy.gemini.ophooks import BaseOpHook

17
colossalai/zero/low_level/low_level_optim.py
View File

@@ -6,6 +6,7 @@ from typing import Dict, Iterator, Optional, Tuple
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed import ProcessGroup
from torch.optim import Optimizer
@@ -617,3 +618,19 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
ret_block_size += current_block_size
yield ret_block, ret_block_size
def update_master_params(self, model: nn.Module) -> None:
"""Update master params from working params
Args:
model (nn.Module): The model to update master params
"""
for p in model.parameters():
p_id = id(p)
if p_id in self._param_store.working_to_master_param:
master_param = self._param_store.working_to_master_param[p_id]
padding_size = self._param_store.get_param_padding_size(p)
working_param = p.data.view(-1)
if padding_size > 0:
working_param = torch.nn.functional.pad(working_param, [0, padding_size])
master_param.copy_(working_param.chunk(self._world_size)[self._local_rank])