26 lines
806 B
Python
26 lines
806 B
Python
# 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")
|