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
fa43bb21
Commit
fa43bb21
authored
May 13, 2022
by
xyupeng
Committed by
binmakeswell
May 17, 2022
Browse files
[NFC] polish colossalai/builder/pipeline.py code style (#951)
parent
c3e423c8
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
4 deletions
+3
-4
colossalai/builder/pipeline.py
colossalai/builder/pipeline.py
+3
-4
No files found.
colossalai/builder/pipeline.py
View file @
fa43bb21
...
...
@@ -251,8 +251,8 @@ def build_pipeline_model(layers: nn.Sequential, num_chunks: int = 1, verbose: bo
partitions
=
partition_uniform
(
len
(
layers
),
pipeline_parallel_size
,
num_chunks
)
module_list
=
[]
for
start
,
end
in
partitions
[
pipeline_rank
]:
module_list
.
append
(
nn
.
Sequential
(
*
[
nn
.
Identity
()
for
_
in
range
(
start
)],
*
layers
[
start
:
end
],
module_list
.
append
(
nn
.
Sequential
(
*
[
nn
.
Identity
()
for
_
in
range
(
start
)],
*
layers
[
start
:
end
],
*
[
nn
.
Identity
()
for
_
in
range
(
len
(
layers
)
-
end
)]))
if
verbose
:
logger
=
get_dist_logger
()
...
...
@@ -264,4 +264,3 @@ def build_pipeline_model(layers: nn.Sequential, num_chunks: int = 1, verbose: bo
log_str
+=
'
\n
'
.
join
([
str
(
layer
)
for
layer
in
layers
[
start
:
end
]])
+
'
\n
'
logger
.
info
(
log_str
,
ranks
=
[
0
])
return
nn
.
ModuleList
(
module_list
)
if
len
(
module_list
)
>
1
else
module_list
[
0
]
\ No newline at end of file
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