colossalai

1
import torch.nn as nn
2
from transformers.models.opt.modeling_opt import OPTAttention
3

4
from .opt_attn import XOPTAttention
5

6

7
def convert_to_xformer_model(model: nn.Module) -> nn.Module:
8
    for module in model.modules():
9
        if isinstance(module, OPTAttention):
10
            module.__class__ = XOPTAttention
11
    return model
12

13

14
def recover_from_xformer_model(model: nn.Module) -> nn.Module:
15
    for module in model.modules():
16
        if isinstance(module, XOPTAttention):
17
            module.__class__ = OPTAttention
18
    return model
19