Skip to content

Commit 4c95061

Browse files
youssefbenchMartinGleize
youssefbench
authored and
MartinGleize
committed
Add OPT architecture
1 parent cbee43e commit 4c95061

7 files changed

Lines changed: 363 additions & 74 deletions

File tree

‎src/fairseq2/composition/models.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,13 @@
3939
register_mistral_configs,
4040
)
4141
from fairseq2.models.nllb import register_nllb_configs
42+
from fairseq2.models.opt import (
43+
OPT_MODEL_FAMILY,
44+
OPTConfig,
45+
convert_opt_state_dict,
46+
create_opt_model,
47+
register_opt_configs,
48+
)
4249
from fairseq2.models.qwen import (
4350
QWEN_FAMILY,
4451
QwenConfig,
@@ -142,6 +149,19 @@ def _register_model_families(container: DependencyContainer) -> None:
142149
)
143150

144151
register_mistral_configs(container)
152+
153+
# OPT
154+
register_model_family(
155+
container,
156+
OPT_MODEL_FAMILY,
157+
kls=TransformerLM,
158+
config_kls=OPTConfig,
159+
factory=create_opt_model,
160+
state_dict_converter=convert_opt_state_dict,
161+
compiler=compile_transformer_lm,
162+
)
163+
164+
register_opt_configs(container)
145165

146166
# NLLB
147167
register_nllb_configs(container)

‎src/fairseq2/models/opt/__init__.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from __future__ import annotations
8+
9+
from fairseq2.models.opt.interop import (
10+
convert_opt_state_dict as convert_opt_state_dict,
11+
)
12+
from fairseq2.models.opt.config import OPT_MODEL_FAMILY as OPT_MODEL_FAMILY
13+
from fairseq2.models.opt.config import OPTConfig as OPTConfig
14+
from fairseq2.models.opt.config import (
15+
register_opt_configs as register_opt_configs,
16+
)
17+
from fairseq2.models.opt.factory import OPTFactory as OPTFactory
18+
from fairseq2.models.opt.factory import (
19+
create_opt_model as create_opt_model,
20+
)
21+
from fairseq2.models.opt.hub import get_opt_model_hub as get_opt_model_hub

‎src/fairseq2/models/opt/config.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from __future__ import annotations
8+
9+
from dataclasses import dataclass
10+
from typing import Final
11+
12+
from fairseq2.runtime.config_registry import ConfigRegistrar
13+
from fairseq2.runtime.dependency import DependencyContainer
14+
15+
OPT_MODEL_FAMILY: Final = "opt"
16+
17+
18+
@dataclass(kw_only=True)
19+
class OPTConfig:
20+
"""Holds the configuration of a OPT model.
21+
22+
The default values correspond to the base architecture as described in
23+
:cite:t:`https://arxiv.org/abs/2205.01068`.
24+
"""
25+
26+
model_dim: int = 768
27+
"""The dimensionality of the model."""
28+
29+
max_seq_len: int = 2048 + 1
30+
"""The maximum sequence length."""
31+
32+
vocab_size: int = 50272
33+
"""The size of the vocabulary."""
34+
35+
pad_idx: int | None = 1
36+
"""The index of the PAD symbol in the vocabulary."""
37+
38+
attn_window_len: int = 2048
39+
"""The local attention window length."""
40+
41+
num_layers: int = 12
42+
"""The number of decoder layers."""
43+
44+
num_attn_heads: int = 12
45+
"""The number of attention heads in decoder layers."""
46+
47+
# If ``None`` or set to ``num_heads``, it is equivalent to standard Multi Head Attention (MHA);
48+
# if set to 1, it is equivalent to Multi Query Attention (MQA).
49+
num_key_value_heads: int = 12
50+
"""The number of key/value heads for Grouped Query Attention."""
51+
52+
ffn_inner_dim: int = 3072
53+
"""The dimensionality of inner projection layers in feed-forward networks."""
54+
55+
dropout_p: float = 0.1
56+
"""The dropout probability on outputs of Transformer layers."""
57+
58+
59+
def register_opt_configs(container: DependencyContainer) -> None:
60+
arch = ConfigRegistrar(container, OPTConfig)
61+
62+
@arch("opt_125m")
63+
def opt_125m() -> OPTConfig:
64+
return OPTConfig()

‎src/fairseq2/models/opt/factory.py‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from __future__ import annotations
8+
9+
from fairseq2.models.transformer import (
10+
CausalAttentionBias,
11+
FeedForwardNetwork,
12+
LocalAttentionStateFactory,
13+
MultiheadAttention,
14+
StandardFeedForwardNetwork,
15+
StandardMultiheadAttention,
16+
TransformerEmbeddingFrontend,
17+
TransformerFrontend,
18+
TransformerNormOrder,
19+
create_default_sdpa,
20+
init_transformer_final_projection,
21+
)
22+
from fairseq2.models.transformer_lm import (
23+
StandardTransformerLMDecoder,
24+
StandardTransformerLMDecoderLayer,
25+
TransformerLM,
26+
TransformerLMDecoder,
27+
TransformerLMDecoderLayer,
28+
)
29+
from fairseq2.nn import (
30+
Embedding,
31+
LayerNorm,
32+
LearnedPositionEncoder,
33+
Linear,
34+
PositionEncoder,
35+
Projection,
36+
RMSNorm,
37+
StandardEmbedding,
38+
StandardLayerNorm,
39+
)
40+
41+
# isort: split
42+
43+
from fairseq2.models.opt.config import OPTConfig
44+
45+
46+
def create_opt_model(config: OPTConfig) -> TransformerLM:
47+
return OPTFactory(config).create_model()
48+
49+
50+
class OPTFactory:
51+
_config: OPTConfig
52+
53+
def __init__(self, config: OPTConfig) -> None:
54+
self._config = config
55+
56+
def create_model(self) -> TransformerLM:
57+
config = self._config
58+
59+
decoder_frontend = self.create_decoder_frontend()
60+
61+
decoder = self.create_decoder()
62+
63+
final_proj = self.create_final_projection()
64+
65+
return TransformerLM(
66+
config.model_dim,
67+
decoder_frontend,
68+
decoder,
69+
final_proj,
70+
config.pad_idx,
71+
config.max_seq_len,
72+
)
73+
74+
def create_decoder_frontend(self) -> TransformerFrontend:
75+
config = self._config
76+
77+
embed = self.create_embedding()
78+
79+
pos_encoder = self.create_position_encoder()
80+
81+
return TransformerEmbeddingFrontend(
82+
config.model_dim,
83+
embed,
84+
pos_encoder=pos_encoder,
85+
no_scale=True,
86+
# dropout_p=config.dropout_p, # TODO: check if there is dropout here
87+
)
88+
89+
def create_embedding(self) -> Embedding:
90+
config = self._config
91+
92+
return StandardEmbedding(config.vocab_size, config.model_dim, config.pad_idx)
93+
94+
def create_decoder(self) -> TransformerLMDecoder:
95+
config = self._config
96+
97+
layers = []
98+
99+
for _ in range(config.num_layers):
100+
layer = self.create_decoder_layer()
101+
102+
layers.append(layer)
103+
104+
layer_norm = self.create_layer_norm()
105+
106+
return StandardTransformerLMDecoder(layers, layer_norm)
107+
108+
def create_position_encoder(self) -> PositionEncoder:
109+
config = self._config
110+
111+
# TODO: should be "= config.model_dim", but fs2 throws error
112+
# See https://github.com/vllm-project/vllm/blob/4719460644b4629db2b6dbf12be331d0b34b4b6f/vllm/model_executor/models/opt.py#L211
113+
encoding_dim = config.model_dim # // config.num_attn_heads
114+
115+
return LearnedPositionEncoder(encoding_dim, config.max_seq_len, offset=2)
116+
117+
def create_decoder_layer(self) -> TransformerLMDecoderLayer:
118+
config = self._config
119+
120+
self_attn = self.create_self_attention()
121+
122+
self_attn_layer_norm = self.create_layer_norm()
123+
124+
ffn = self.create_ffn()
125+
126+
ffn_layer_norm = self.create_layer_norm()
127+
128+
return StandardTransformerLMDecoderLayer(
129+
self_attn,
130+
self_attn_layer_norm,
131+
ffn,
132+
ffn_layer_norm,
133+
norm_order=TransformerNormOrder.PRE,
134+
dropout_p=config.dropout_p,
135+
)
136+
137+
def create_self_attention(self) -> MultiheadAttention:
138+
config = self._config
139+
140+
attn_bias = CausalAttentionBias(attn_window_len=config.attn_window_len)
141+
142+
sdpa = create_default_sdpa(attn_bias)
143+
144+
incremental_state_factory = LocalAttentionStateFactory(config.attn_window_len)
145+
146+
return StandardMultiheadAttention(
147+
config.model_dim,
148+
config.num_attn_heads,
149+
sdpa,
150+
num_key_value_heads=config.num_key_value_heads,
151+
# pos_encoder=pos_encoder,
152+
bias=True,
153+
state_factory=incremental_state_factory,
154+
)
155+
156+
def create_ffn(self) -> FeedForwardNetwork:
157+
config = self._config
158+
159+
return StandardFeedForwardNetwork(config.model_dim, config.ffn_inner_dim, bias=True)
160+
161+
def create_layer_norm(self) -> LayerNorm:
162+
config = self._config
163+
164+
return StandardLayerNorm(config.model_dim, bias=True)
165+
166+
def create_final_projection(self) -> Projection:
167+
config = self._config
168+
169+
return Linear(
170+
config.model_dim,
171+
config.vocab_size,
172+
bias=False,
173+
init_fn=init_transformer_final_projection,
174+
)

‎src/fairseq2/models/opt/hub.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from __future__ import annotations
8+
9+
from fairseq2.models import ModelHubAccessor
10+
from fairseq2.models.transformer_lm import TransformerLM
11+
12+
# isort: split
13+
14+
from fairseq2.models.opt.config import OPTConfig, OPT_MODEL_FAMILY
15+
16+
get_opt_model_hub = ModelHubAccessor(OPT_MODEL_FAMILY, TransformerLM, OPTConfig)

‎src/fairseq2/models/opt/interop.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from __future__ import annotations
8+
9+
from fairseq2.models.utils.checkpoint import convert_state_dict
10+
11+
# isort: split
12+
13+
from fairseq2.models.opt.config import OPTConfig
14+
15+
16+
_OPT_HG_KEY_MAP = {
17+
# fmt: off
18+
r"^model\.decoder\.embed_tokens\.": r"decoder_frontend.embed.",
19+
r"^model\.decoder\.embed_positions\.": r"decoder_frontend.pos_encoder.",
20+
r"^model\.decoder\.layers\.([0-9]+)\.self_attn_layer_norm\.": r"decoder.layers.\1.self_attn_layer_norm.",
21+
r"^model\.decoder\.layers\.([0-9]+)\.self_attn\.q_proj\.": r"decoder.layers.\1.self_attn.q_proj.",
22+
r"^model\.decoder\.layers\.([0-9]+)\.self_attn\.k_proj\.": r"decoder.layers.\1.self_attn.k_proj.",
23+
r"^model\.decoder\.layers\.([0-9]+)\.self_attn\.v_proj\.": r"decoder.layers.\1.self_attn.v_proj.",
24+
r"^model\.decoder\.layers\.([0-9]+)\.self_attn\.out_proj\.": r"decoder.layers.\1.self_attn.output_proj.",
25+
r"^model\.decoder\.layers\.([0-9]+)\.mlp\.gate_proj\.": r"decoder.layers.\1.ffn.gate_proj.",
26+
r"^model\.decoder\.layers\.([0-9]+)\.fc1\.": r"decoder.layers.\1.ffn.inner_proj.",
27+
r"^model\.decoder\.layers\.([0-9]+)\.fc2\.": r"decoder.layers.\1.ffn.output_proj.",
28+
r"^model\.decoder\.layers\.([0-9]+)\.final_layer_norm\.": r"decoder.layers.\1.ffn_layer_norm.",
29+
r"^model\.decoder\.final_layer_norm\.": r"decoder.layer_norm.",
30+
r"^lm_head\.": r"final_proj.",
31+
# fmt: on
32+
}
33+
34+
35+
def convert_opt_state_dict(
36+
state_dict: dict[str, object], config: OPTConfig
37+
) -> dict[str, object]:
38+
if "model.decoder.embed_tokens.weight" in state_dict: # Hugging Face
39+
state_dict = convert_state_dict(state_dict, _OPT_HG_KEY_MAP)
40+
41+
return state_dict

0 commit comments

Comments
 (0)