mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-03 05:58:09 +00:00
19 lines
533 B
Python
19 lines
533 B
Python
import torch.nn as nn
|
|
from transformers.models.opt.modeling_opt import OPTAttention
|
|
|
|
from .opt_attn import XOPTAttention
|
|
|
|
|
|
def convert_to_xformer_model(model: nn.Module) -> nn.Module:
|
|
for module in model.modules():
|
|
if isinstance(module, OPTAttention):
|
|
module.__class__ = XOPTAttention
|
|
return model
|
|
|
|
|
|
def recover_from_xformer_model(model: nn.Module) -> nn.Module:
|
|
for module in model.modules():
|
|
if isinstance(module, XOPTAttention):
|
|
module.__class__ = OPTAttention
|
|
return model
|