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

[hotfix] Fix examples no pad token & auto parallel codegen bug; (#5606)

Browse Source 
* fix no pad token bug

* fixed some auto parallel codegen bug, but might not run on torch 2.1

---------

Co-authored-by: Edenzzzz <wtan45@wisc.edu>
This commit is contained in:
Edenzzzz
2024-04-18 18:15:50 +08:00
committed by GitHub
parent a0ad587c24
commit d83c633ca6
6 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/_analyzer/fx/codegen.py
View File

@@ -246,7 +246,7 @@ def emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func,
@compatibility(is_backward_compatible=True)
class ActivationCheckpointCodeGen(CodeGen):
def _gen_python_code(self, nodes, root_module: str, namespace: _Namespace) -> PythonCode:
def _gen_python_code(self, nodes, root_module: str, namespace: _Namespace, verbose=None) -> PythonCode:
free_vars: List[str] = []
body: List[str] = []
globals_: Dict[str, Any] = {}

2
colossalai/auto_parallel/offload/base_offload_module.py
View File

@@ -5,7 +5,7 @@ import torch
import torch.nn as nn
from colossalai.utils import _cast_float
from colossalai.zero.legacy.gemini.tensor_utils import free_storage
from colossalai.utils.common import free_storage
from .region_manager import RegionManager
from .util import GlobalRuntimeInfo

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

@@ -3,7 +3,8 @@ from typing import Dict, List, Tuple
import torch
from torch.fx import Node
from colossalai.zero.legacy.gemini.tensor_utils import alloc_storage, free_storage
from colossalai.utils.common import free_storage
from colossalai.zero.gemini.chunk.chunk import alloc_storage
class Region:

2
colossalai/autochunk/autochunk_codegen.py
View File

@@ -372,7 +372,7 @@ if AUTOCHUNK_AVAILABLE:
if print_progress:
get_logger().info("AutoChunk start codegen")
def _gen_python_code(self, nodes, root_module: str, namespace: _Namespace) -> PythonCode:
def _gen_python_code(self, nodes, root_module: str, namespace: _Namespace, verbose=None) -> PythonCode:
free_vars: List[str] = []
body: List[str] = []
globals_: Dict[str, Any] = {}

2
colossalai/fx/codegen/activation_checkpoint_codegen.py
View File

@@ -625,7 +625,7 @@ def emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func,
if CODEGEN_AVAILABLE:
class ActivationCheckpointCodeGen(CodeGen):
def _gen_python_code(self, nodes, root_module: str, namespace: _Namespace) -> PythonCode:
def _gen_python_code(self, nodes, root_module: str, namespace: _Namespace, verbose=None) -> PythonCode:
free_vars: List[str] = []
body: List[str] = []
globals_: Dict[str, Any] = {}