Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/mmcv.git" did not exist on "e22740b1d6953d75a0acecce4455d23800b1f018"
Commit a8c5a3cb authored by Rick Ho's avatar Rick Ho
Browse files

bug fix for megatron patch

parent ae10e942
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
fmoe/megatron/patch.py
View file @ a8c5a3cb
r""" r"""
Patching some of Megatron-LM's functions to create an MoE model Patching some of Megatron-LM's functions to create an MoE model
""" """
import torch
def patch_forward_step(forward_step_func): def patch_forward_step(forward_step_func):
r""" r"""
Patch model's forward_step_func to support balance loss Patch model's forward_step_func to support balance loss
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment