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
evt_fugx1
dcu_megatron
Commits
770fa304
Commit
770fa304
authored
Apr 25, 2025
by
dongcl
Browse files
修改mtp
parent
8096abd4
Changes
44
Expand all
Hide whitespace changes
Inline
Side-by-side
Showing
4 changed files
with
106 additions
and
140 deletions
+106
-140
dcu_megatron/core/utils.py
dcu_megatron/core/utils.py
+0
-30
dcu_megatron/training/arguments.py
dcu_megatron/training/arguments.py
+10
-8
dcu_megatron/training/utils.py
dcu_megatron/training/utils.py
+82
-88
pretrain_gpt.py
pretrain_gpt.py
+14
-14
No files found.
dcu_megatron/core/utils.py
View file @
770fa304
...
...
@@ -30,33 +30,3 @@ def is_flux_min_version(version, check_equality=True):
if
check_equality
:
return
get_flux_version
()
>=
PkgVersion
(
version
)
return
get_flux_version
()
>
PkgVersion
(
version
)
def
tensor_slide
(
tensor
:
Optional
[
torch
.
Tensor
],
num_slice
:
int
,
dims
:
Union
[
int
,
List
[
int
]]
=
-
1
,
step
:
int
=
1
,
return_first
=
False
,
)
->
List
[
Union
[
torch
.
Tensor
,
None
]]:
"""通用滑动窗口函数,支持任意维度"""
if
tensor
is
None
:
# return `List[None]` to avoid NoneType Error
return
[
None
]
*
(
num_slice
+
1
)
if
num_slice
==
0
:
return
[
tensor
]
window_size
=
tensor
.
shape
[
-
1
]
-
num_slice
dims
=
[
dims
]
if
isinstance
(
dims
,
int
)
else
sorted
(
dims
,
reverse
=
True
)
# 连续多维度滑动
slices
=
[]
for
i
in
range
(
0
,
tensor
.
size
(
dims
[
-
1
])
-
window_size
+
1
,
step
):
slice_obj
=
[
slice
(
None
)]
*
tensor
.
dim
()
for
dim
in
dims
:
slice_obj
[
dim
]
=
slice
(
i
,
i
+
window_size
)
slices
.
append
(
tensor
[
tuple
(
slice_obj
)])
if
return_first
:
return
slices
return
slices
dcu_megatron/training/arguments.py
View file @
770fa304
This diff is collapsed.
Click to expand it.
dcu_megatron/training/utils.py
View file @
770fa304
This diff is collapsed.
Click to expand it.
pretrain_gpt.py
View file @
770fa304
This diff is collapsed.
Click to expand it.
Prev
1
2
3
Next
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