初始化项目,由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
25
slime/utils/megatron_bridge_utils.py
Normal file
25
slime/utils/megatron_bridge_utils.py
Normal file
@@ -0,0 +1,25 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
try:
|
||||
from megatron.core.utils import unwrap_model
|
||||
except ImportError:
|
||||
unwrap_model = None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def patch_megatron_model(model):
|
||||
unwrapped_model = unwrap_model(model)[0]
|
||||
model_config = unwrapped_model.config
|
||||
attribute_was_added = False
|
||||
if not hasattr(model_config, "share_embeddings_and_output_weights"):
|
||||
model_config.share_embeddings_and_output_weights = unwrapped_model.share_embeddings_and_output_weights
|
||||
attribute_was_added = True
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if attribute_was_added:
|
||||
delattr(model_config, "share_embeddings_and_output_weights")
|
||||
Reference in New Issue
Block a user