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

[zero] reorganize zero/gemini folder structure (#3424)

Browse Source 
* [zero] refactor low-level zero folder structure

* [zero] fix legacy zero import path

* [zero] fix legacy zero import path

* [zero] remove useless import

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor legacy zero import path

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor legacy zero import path

* [zero] fix test import path

* [zero] fix test

* [zero] fix circular import

* [zero] update import
This commit is contained in:
ver217
2023-04-04 13:48:16 +08:00
committed by GitHub
parent b09adff724
commit 26b7aac0be
142 changed files with 1435 additions and 1404 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/auto_parallel/offload/region.py
View File

@@ -1,7 +1,10 @@
from typing import List, Dict, Tuple
from typing import Dict, List, Tuple
import torch
from torch.fx import Node
from colossalai.gemini.tensor_utils import alloc_storage, free_storage
from colossalai.zero.legacy.gemini.tensor_utils import alloc_storage, free_storage
class Region:
"""
@@ -52,15 +55,13 @@ class Region:
Map the parameters in the region to a contiguous memory space.
"""
self.fp16_data = torch.zeros(
self.param_num, dtype=torch.half, device='cuda')
self.fp16_data = torch.zeros(self.param_num, dtype=torch.half, device='cuda')
offset = 0
for param in self.fp16_params:
param.data = param.data.cuda()
p_num = param.data.numel()
self.fp16_data[offset:offset + p_num].copy_(param.data.flatten())
param.data = self.fp16_data[offset:offset +
p_num].view(param.data.shape)
param.data = self.fp16_data[offset:offset + p_num].view(param.data.shape)
self.param_to_range[param] = (offset, offset + p_num)
offset += p_num
@@ -141,4 +142,4 @@ class Region:
def __update_params_ptr(self) -> None:
for param in self.fp16_params:
begin, end = self.param_to_range[param]
param.data = self.fp16_data[begin:end].view(param.data.shape)
param.data = self.fp16_data[begin:end].view(param.data.shape)