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

[autochunk] support transformer (#2526)

Browse Source
This commit is contained in:
oahzxl
2023-01-31 16:00:06 +08:00
committed by GitHub
parent 6e0faa70e0
commit 63199c6687
20 changed files with 1214 additions and 1084 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/autochunk/autochunk_codegen.py
View File

@@ -3,9 +3,12 @@ from typing import Any, Dict, Iterable, List, Tuple
import torch
import colossalai
from colossalai.fx._compatibility import is_compatible_with_meta
from colossalai.fx.codegen.activation_checkpoint_codegen import CODEGEN_AVAILABLE
if CODEGEN_AVAILABLE:
AUTOCHUNK_AVAILABLE = CODEGEN_AVAILABLE and is_compatible_with_meta()
if AUTOCHUNK_AVAILABLE:
from torch.fx.graph import (
CodeGen,
PythonCode,
@@ -272,7 +275,7 @@ def emit_code_with_chunk(
node_idx += 1
if CODEGEN_AVAILABLE:
if AUTOCHUNK_AVAILABLE:
class AutoChunkCodeGen(CodeGen):