Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
c7c224ee
Commit
c7c224ee
authored
Apr 02, 2022
by
Ziheng Qin
Committed by
binmakeswell
Apr 06, 2022
Browse files
[NFC] polish colossalai/builder/pipeline.py code style (#638)
parent
10591ecd
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
6 additions
and
3 deletions
+6
-3
colossalai/builder/pipeline.py
colossalai/builder/pipeline.py
+6
-3
No files found.
colossalai/builder/pipeline.py
View file @
c7c224ee
import
copy
import
copy
import
heapq
import
heapq
from
colossalai.builder
import
build_model
,
build_layer
from
colossalai.builder
import
build_model
,
build_layer
from
colossalai.context.parallel_mode
import
ParallelMode
from
colossalai.context.parallel_mode
import
ParallelMode
from
colossalai.core
import
global_context
as
gpc
from
colossalai.core
import
global_context
as
gpc
...
@@ -40,6 +39,7 @@ def _binary_partition(weights, st, ed):
...
@@ -40,6 +39,7 @@ def _binary_partition(weights, st, ed):
def
_heap_addition
(
weights
,
intervals
,
add_cnt
):
def
_heap_addition
(
weights
,
intervals
,
add_cnt
):
"""
"""
"""
"""
def
_heap_push
(
heap
,
st
,
ed
):
def
_heap_push
(
heap
,
st
,
ed
):
value
=
weights
[
ed
-
1
]
value
=
weights
[
ed
-
1
]
if
st
>
0
:
if
st
>
0
:
...
@@ -162,7 +162,10 @@ def count_layer_params(layers):
...
@@ -162,7 +162,10 @@ def count_layer_params(layers):
return
param_counts
return
param_counts
def
build_pipeline_model_from_cfg
(
config
,
num_chunks
:
int
=
1
,
partition_method
:
str
=
'parameter'
,
verbose
:
bool
=
False
):
def
build_pipeline_model_from_cfg
(
config
,
num_chunks
:
int
=
1
,
partition_method
:
str
=
'parameter'
,
verbose
:
bool
=
False
):
"""An initializer to split the model into different stages for pipeline parallelism.
"""An initializer to split the model into different stages for pipeline parallelism.
An example for the model config is shown below. The class VisionTransformerFromConfig should
An example for the model config is shown below. The class VisionTransformerFromConfig should
...
@@ -218,7 +221,7 @@ def build_pipeline_model_from_cfg(config, num_chunks: int = 1, partition_method:
...
@@ -218,7 +221,7 @@ def build_pipeline_model_from_cfg(config, num_chunks: int = 1, partition_method:
log_str
+=
f
'
\n
===== stage=
{
stage
}
, layers=
{
num_layers
}
=====
\n
'
log_str
+=
f
'
\n
===== stage=
{
stage
}
, layers=
{
num_layers
}
=====
\n
'
for
st
,
ed
in
parts
[
stage
]:
for
st
,
ed
in
parts
[
stage
]:
for
idx
,
layer
in
enumerate
(
layers
[
st
:
ed
]):
for
idx
,
layer
in
enumerate
(
layers
[
st
:
ed
]):
log_str
+=
f
'
\t
{
idx
+
st
:
2
d
}
:
{
layer
}
\n
'
log_str
+=
f
'
\t
{
idx
+
st
:
2
d
}
:
{
layer
}
\n
'
logger
.
info
(
log_str
,
ranks
=
[
0
])
logger
.
info
(
log_str
,
ranks
=
[
0
])
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment