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

[shardformer] add gpt2 policy and modify shard and slicer to support (#3883)

Browse Source 
* add gpt2 policy and modify shard and slicer to support

* remove unused code

* polish code
This commit is contained in:
FoolPlayer
2023-06-07 16:09:40 +08:00
committed by Frank Lee
parent 70173e3123
commit 79f8d5d54b
7 changed files with 233 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/shardformer/policies/bert.py
View File

@@ -1,4 +1,3 @@
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Tuple, Type
import torch.nn as nn