@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 085ac221a3

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 12:15 UTC

Deprecate `JAISLMHeadModel` (#43784)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 28, 2026 at 19:29 UTC 085ac221a32bca4d9ef4d325444b9713412a7ddf

8 files changed +1 -651

docs/models/supported_models.md

-1

@@ -428,7 +428,6 @@ th {

428 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |

432 | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |

tests/distributed/test_pipeline_parallel.py

-1

@@ -127,7 +127,6 @@ TEXT_GENERATION_MODELS = {

127 # Uses Llama

128 # "internlm/internlm-chat-7b": PPTestSettings.fast(),

129 "internlm/internlm2-chat-7b": PPTestSettings.fast(),

130 - "inceptionai/jais-13b-chat": PPTestSettings.fast(),

130 "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),

131 "pfnet/plamo-2-1b": PPTestSettings.fast(),

132 "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(),

tests/models/registry.py

-1

@@ -363,7 +363,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

363 "IQuestLoopCoderForCausalLM": _HfExamplesInfo(

364 "IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct", trust_remote_code=True

365 ),

366 - "JAISLMHeadModel": _HfExamplesInfo("inceptionai/jais-13b-chat"),

366 "Jais2ForCausalLM": _HfExamplesInfo(

367 "inceptionai/Jais-2-8B-Chat", min_transformers_version="4.58"

368 ),

vllm/model_executor/models/jais.py deleted

-401

@@ -1,401 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -# Adapted from

5 -# https://huggingface.co/inceptionai/jais-30b-chat-v3/blob/main/modeling_jais.py

8 -# reserved.

10 -#

11 -# Licensed under the Apache License, Version 2.0 (the "License");

12 -# you may not use this file except in compliance with the License.

13 -# You may obtain a copy of the License at

14 -#

15 -# http://www.apache.org/licenses/LICENSE-2.0

16 -#

17 -# Unless required by applicable law or agreed to in writing, software

18 -# distributed under the License is distributed on an "AS IS" BASIS,

19 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

20 -# See the License for the specific language governing permissions and

21 -# limitations under the License.

22 -"""Inference-only Jais model compatible with HuggingFace weights."""

23 -

24 -import math

25 -from collections.abc import Iterable

26 -from itertools import islice

27 -

28 -import torch

29 -from torch import nn

30 -

31 -from vllm.compilation.decorators import support_torch_compile

32 -from vllm.config import CacheConfig, VllmConfig

33 -from vllm.distributed import (

34 - get_pp_group,

35 - get_tensor_model_parallel_rank,

36 - get_tensor_model_parallel_world_size,

37 -)

38 -from vllm.model_executor.layers.attention import Attention

39 -from vllm.model_executor.layers.linear import (

40 - ColumnParallelLinear,

41 - QKVParallelLinear,

42 - RowParallelLinear,

43 -)

44 -from vllm.model_executor.layers.logits_processor import LogitsProcessor

45 -from vllm.model_executor.layers.quantization import QuantizationConfig

46 -from vllm.model_executor.layers.vocab_parallel_embedding import (

47 - ParallelLMHead,

48 - VocabParallelEmbedding,

49 -)

50 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

51 -from vllm.sequence import IntermediateTensors

52 -from vllm.transformers_utils.configs.jais import JAISConfig

53 -

54 -from .interfaces import SupportsPP

55 -from .utils import (

56 - AutoWeightsLoader,

57 - is_pp_missing_parameter,

58 - make_empty_intermediate_tensors_factory,

59 - make_layers,

60 - maybe_prefix,

61 -)

62 -

63 -

64 -class SwiGLUActivation(nn.Module):

65 - def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:

66 - return x1 * nn.functional.silu(x2)

67 -

68 -

69 -def _get_alibi_slopes(n):

70 - def get_slopes_power_of_2(n):

71 - start = 2 ** (-(2 ** -(math.log2(n) - 3)))

72 - ratio = start

73 - return [start * ratio**i for i in range(n)]

74 -

75 - if math.log2(n).is_integer():

76 - return get_slopes_power_of_2(n)

77 - else:

78 - closest_power_of_2 = 2 ** math.floor(math.log2(n))

79 - return (

80 - get_slopes_power_of_2(closest_power_of_2)

81 - + _get_alibi_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]

82 - )

83 -

84 -

85 -class JAISAttention(nn.Module):

86 - def __init__(

87 - self,

88 - config: JAISConfig,

89 - cache_config: CacheConfig | None = None,

90 - quant_config: QuantizationConfig | None = None,

91 - prefix: str = "",

92 - ):

93 - super().__init__()

94 - self.hidden_size = config.hidden_size

95 - total_num_heads = config.num_attention_heads

96 - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()

97 - assert total_num_heads % tensor_model_parallel_world_size == 0

98 - self.num_heads = total_num_heads // tensor_model_parallel_world_size

99 - self.head_dim = self.hidden_size // total_num_heads

100 - if hasattr(config, "scale_qk_dot_by_d"):

101 - config.mup_scale_qk_dot_by_d = config.scale_qk_dot_by_d

102 - self.attn_scale_power = 1.0 if config.mup_scale_qk_dot_by_d else 0.5

103 - self.scale = self.head_dim**-self.attn_scale_power

104 -

105 - self.c_attn = QKVParallelLinear(

106 - self.hidden_size,

107 - self.head_dim,

108 - total_num_heads,

109 - bias=True,

110 - quant_config=quant_config,

111 - prefix=f"{prefix}.c_attn",

112 - )

113 - self.c_proj = RowParallelLinear(

114 - self.hidden_size,

115 - self.hidden_size,

116 - bias=True,

117 - quant_config=quant_config,

118 - prefix=f"{prefix}.c_proj",

119 - )

120 -

121 - self.use_alibi = config.position_embedding_type == "alibi"

122 - alibi_slopes = None

123 - if self.use_alibi:

124 - tp_rank = get_tensor_model_parallel_rank()

125 - head_start = tp_rank * self.num_heads

126 - head_end = (tp_rank + 1) * self.num_heads

127 - alibi_slopes = _get_alibi_slopes(total_num_heads)

128 - alibi_slopes = alibi_slopes[head_start:head_end]

129 - self.attn = Attention(

130 - self.num_heads,

131 - self.head_dim,

132 - scale=self.scale,

133 - alibi_slopes=alibi_slopes,

134 - cache_config=cache_config,

135 - quant_config=quant_config,

136 - prefix=f"{prefix}.attn",

137 - )

138 -

139 - def forward(

140 - self,

141 - hidden_states: torch.Tensor,

142 - ) -> torch.Tensor:

143 - qkv, _ = self.c_attn(hidden_states)

144 - q, k, v = qkv.chunk(chunks=3, dim=-1)

145 - attn_output = self.attn(q, k, v)

146 - attn_output, _ = self.c_proj(attn_output)

147 - return attn_output

148 -

149 -

150 -class JAISMLP(nn.Module):

151 - def __init__(

152 - self,

153 - intermediate_size: int,

154 - config: JAISConfig,

155 - quant_config: QuantizationConfig | None = None,

156 - prefix: str = "",

157 - ):

158 - super().__init__()

159 - hidden_size = config.hidden_size

160 - self.swiglu = config.activation_function == "swiglu"

161 - self.c_fc = ColumnParallelLinear(

162 - hidden_size,

163 - intermediate_size,

164 - bias=True,

165 - quant_config=quant_config,

166 - prefix=f"{prefix}.c_fc",

167 - )

168 - self.c_fc2 = (

169 - ColumnParallelLinear(

170 - hidden_size,

171 - intermediate_size,

172 - bias=True,

173 - quant_config=quant_config,

174 - prefix=f"{prefix}.c_fc2",

175 - )

176 - if self.swiglu

177 - else None

178 - )

179 - self.c_proj = RowParallelLinear(

180 - intermediate_size,

181 - hidden_size,

182 - bias=True,

183 - quant_config=quant_config,

184 - prefix=f"{prefix}.c_proj",

185 - )

186 -

187 - self.act = SwiGLUActivation()

188 -

189 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

190 - if self.swiglu:

191 - hidden_states2, _ = self.c_fc2(hidden_states)

192 - hidden_states, _ = self.c_fc(hidden_states)

193 - hidden_states = (

194 - self.act(hidden_states, hidden_states2)

195 - if self.swiglu

196 - else self.act(hidden_states)

197 - )

198 - hidden_states, _ = self.c_proj(hidden_states)

199 - return hidden_states

200 -

201 -

202 -class JAISBlock(nn.Module):

203 - def __init__(

204 - self,

205 - config: JAISConfig,

206 - cache_config: CacheConfig | None = None,

207 - quant_config: QuantizationConfig | None = None,

208 - prefix: str = "",

209 - ):

210 - super().__init__()

211 - hidden_size = config.hidden_size

212 - inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size

213 -

214 - self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)

215 - self.attn = JAISAttention(

216 - config, cache_config, quant_config, prefix=f"{prefix}.attn"

217 - )

218 - self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)

219 - self.mlp = JAISMLP(inner_dim, config, quant_config, prefix=f"{prefix}.mlp")

220 -

221 - def forward(

222 - self,

223 - hidden_states: torch.Tensor,

224 - ) -> torch.Tensor:

225 - residual = hidden_states

226 - hidden_states = self.ln_1(hidden_states)

227 - attn_output = self.attn(

228 - hidden_states=hidden_states,

229 - )

230 - # residual connection

231 - hidden_states = attn_output + residual

232 -

233 - residual = hidden_states

234 - hidden_states = self.ln_2(hidden_states)

235 - feed_forward_hidden_states = self.mlp(hidden_states)

236 - # residual connection

237 - hidden_states = residual + feed_forward_hidden_states

238 - return hidden_states

239 -

240 -

241 -@support_torch_compile

242 -class JAISModel(nn.Module):

243 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

244 - super().__init__()

245 -

246 - config = vllm_config.model_config.hf_config

247 - cache_config = vllm_config.cache_config

248 - quant_config = vllm_config.quant_config

249 -

250 - self.config = config

251 - assert not config.scale_attn_by_inverse_layer_idx

252 - assert not config.reorder_and_upcast_attn

253 - self.embed_dim = config.hidden_size

254 - self.wte = VocabParallelEmbedding(config.vocab_size, self.embed_dim)

255 - self.wpe = (

256 - nn.Embedding(config.max_position_embeddings, self.embed_dim)

257 - if config.position_embedding_type != "alibi"

258 - else None

259 - )

260 - if hasattr(config, "embeddings_scale"):

261 - self.embeddings_scale = config.embeddings_scale

262 - else:

263 - self.embeddings_scale = config.mup_embeddings_scale

264 -

265 - self.start_layer, self.end_layer, self.h = make_layers(

266 - config.num_hidden_layers,

267 - lambda prefix: JAISBlock(

268 - config=config,

269 - cache_config=cache_config,

270 - quant_config=quant_config,

271 - prefix=prefix,

272 - ),

273 - prefix=f"{prefix}.h",

274 - )

275 -

276 - self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)

277 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

278 - ["hidden_states"], config.n_embd

279 - )

280 -

281 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

282 - return self.wte(input_ids)

283 -

284 - def forward(

285 - self,

286 - input_ids: torch.Tensor | None,

287 - position_ids: torch.Tensor,

288 - intermediate_tensors: IntermediateTensors | None = None,

289 - inputs_embeds: torch.Tensor | None = None,

290 - ) -> IntermediateTensors | torch.Tensor:

291 - if get_pp_group().is_first_rank:

292 - if inputs_embeds is None:

293 - inputs_embeds = self.embed_input_ids(input_ids)

294 - if self.wpe is not None:

295 - position_embeds = self.wpe(position_ids)

296 - hidden_states = inputs_embeds + position_embeds

297 - else:

298 - hidden_states = inputs_embeds

299 - hidden_states *= torch.tensor(

300 - float(self.embeddings_scale), dtype=hidden_states.dtype

301 - )

302 - else:

303 - assert intermediate_tensors is not None

304 - hidden_states = intermediate_tensors["hidden_states"]

305 -

306 - for layer in islice(self.h, self.start_layer, self.end_layer):

307 - hidden_states = layer(hidden_states)

308 -

309 - if not get_pp_group().is_last_rank:

310 - return IntermediateTensors({"hidden_states": hidden_states})

311 -

312 - hidden_states = self.ln_f(hidden_states)

313 - return hidden_states

314 -

315 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

316 - params_dict = dict(self.named_parameters(remove_duplicate=False))

317 - loaded_params: set[str] = set()

318 - for name, loaded_weight in weights:

319 - if ".attn.bias" in name or ".attn.masked_bias" in name:

320 - # Skip attention mask.

321 - # NOTE: "c_attn.bias" should not be skipped.

322 - continue

323 - if "relative_pe" in name:

324 - continue

325 -

326 - if is_pp_missing_parameter(name, self):

327 - continue

328 -

329 - param = params_dict[name]

330 - # The HF's GPT-2 implementation uses Conv1D instead of Linear.

331 - # Because of this, we need to transpose the weights.

332 - # Note(zhuohan): the logic below might break quantized models.

333 - for conv1d_weight_name in ["c_attn", "c_proj", "c_fc"]:

334 - if conv1d_weight_name not in name:

335 - continue

336 - if not name.endswith(".weight"):

337 - continue

338 - loaded_weight = loaded_weight.t()

339 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

340 - weight_loader(param, loaded_weight)

341 - loaded_params.add(name)

342 - return loaded_params

343 -

344 -

345 -class JAISLMHeadModel(nn.Module, SupportsPP):

346 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

347 - super().__init__()

348 - config = vllm_config.model_config.hf_config

349 - quant_config = vllm_config.quant_config

350 - self.config = config

351 - self.quant_config = quant_config

352 - self.transformer = JAISModel(

353 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer")

354 - )

355 - if self.config.tie_word_embeddings:

356 - self.lm_head = self.transformer.wte

357 - else:

358 - self.lm_head = ParallelLMHead(

359 - self.config.vocab_size,

360 - self.config.hidden_size,

361 - prefix=maybe_prefix(prefix, "lm_head"),

362 - )

363 - if hasattr(config, "width_scale"):

364 - self.output_logits_scale = config.width_scale

365 - else:

366 - self.output_logits_scale = config.mup_output_alpha * config.mup_width_scale

367 - self.logits_processor = LogitsProcessor(

368 - vocab_size=config.vocab_size, scale=self.output_logits_scale

369 - )

370 - self.make_empty_intermediate_tensors = (

371 - self.transformer.make_empty_intermediate_tensors

372 - )

373 -

374 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

375 - return self.transformer.embed_input_ids(input_ids)

376 -

377 - def forward(

378 - self,

379 - input_ids: torch.Tensor | None,

380 - positions: torch.Tensor,

381 - intermediate_tensors: IntermediateTensors | None = None,

382 - inputs_embeds: torch.Tensor | None = None,

383 - ) -> IntermediateTensors | torch.Tensor:

384 - hidden_states = self.transformer(

385 - input_ids, positions, intermediate_tensors, inputs_embeds

386 - )

387 - return hidden_states

388 -

389 - def compute_logits(

390 - self,

391 - hidden_states: torch.Tensor,

392 - ) -> torch.Tensor | None:

393 - logits = self.logits_processor(self.lm_head, hidden_states)

394 - return logits

395 -

396 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

397 - loader = AutoWeightsLoader(

398 - self,

399 - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

400 - )

401 - return loader.load_weights(weights)

vllm/model_executor/models/registry.py

+1 -1

@@ -146,7 +146,6 @@ _TEXT_GENERATION_MODELS = {

146 "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),

147 "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),

148 "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),

149 - "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),

149 "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),

150 "JambaForCausalLM": ("jamba", "JambaForCausalLM"),

151 "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),

@@ -705,6 +704,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = {

704 "Phi3SmallForCausalLM": "0.9.2",

705 "Phi4FlashForCausalLM": "0.10.2",

706 "Phi4MultimodalForCausalLM": "0.12.0",

707 + "JAISLMHeadModel": "0.22.0",

708 # encoder-decoder models except whisper

709 # have been removed for V0 deprecation.

710 "DonutForConditionalGeneration": "0.10.2",

vllm/transformers_utils/config.py

-1

@@ -113,7 +113,6 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

113 kimi_k25="KimiK25Config",

114 RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)

115 RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct)

116 - jais="JAISConfig",

116 mlp_speculator="MLPSpeculatorConfig",

117 medusa="MedusaConfig",

118 midashenglm="MiDashengLMConfig",

vllm/transformers_utils/configs/__init__.py

-2

@@ -46,7 +46,6 @@ _CLASS_TO_MODULE: dict[str, str] = {

46 # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the

47 # `FalconConfig` class from the official HuggingFace transformers library.

48 "RWConfig": "vllm.transformers_utils.configs.falcon",

49 - "JAISConfig": "vllm.transformers_utils.configs.jais",

49 "LagunaConfig": "vllm.transformers_utils.configs.laguna",

50 "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",

51 "MedusaConfig": "vllm.transformers_utils.configs.medusa",

@@ -115,7 +114,6 @@ __all__ = [

114 "HyperCLOVAXConfig",

115 "IsaacConfig",

116 "RWConfig",

118 - "JAISConfig",

117 "LagunaConfig",

118 "Lfm2MoeConfig",

119 "MedusaConfig",

vllm/transformers_utils/configs/jais.py deleted

-243

@@ -1,243 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

7 -#

8 -# Licensed under the Apache License, Version 2.0 (the "License");

9 -# you may not use this file except in compliance with the License.

10 -# You may obtain a copy of the License at

11 -#

12 -# http://www.apache.org/licenses/LICENSE-2.0

13 -#

14 -# Unless required by applicable law or agreed to in writing, software

15 -# distributed under the License is distributed on an "AS IS" BASIS,

16 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

17 -# See the License for the specific language governing permissions and

18 -# limitations under the License.

19 -"""JAIS configuration"""

20 -

21 -from transformers.configuration_utils import PretrainedConfig

22 -from transformers.utils import logging

23 -

24 -logger = logging.get_logger(__name__)

25 -

26 -

27 -class JAISConfig(PretrainedConfig):

28 - """

29 - This is the configuration class to store the configuration of a

30 - [`JAISModel`]. It is used to instantiate a JAIS model according to the

31 - specified arguments, defining the model architecture.

32 -

33 - Configuration objects inherit from [`PretrainedConfig`] and can be used

34 - to control the model outputs. Read the documentation from

35 - [`PretrainedConfig`] for more information.

36 -

37 -

38 - Args:

39 - vocab_size (`int`, *optional*, defaults to 50257):

40 - Vocabulary size of the JAIS model. Defines the number of different

41 - tokens that can be represented by the

42 - `inputs_ids` passed when calling [`JAISModel`].

43 - n_positions (`int`, *optional*, defaults to 1024):

44 - The maximum sequence length that this model might ever be used

45 - with. Typically set this to something large just in case

46 - (e.g., 512 or 1024 or 2048).

47 - n_embd (`int`, *optional*, defaults to 768):

48 - Dimensionality of the embeddings and hidden states.

49 - n_layer (`int`, *optional*, defaults to 12):

50 - Number of hidden layers in the Transformer encoder.

51 - n_head (`int`, *optional*, defaults to 12):

52 - Number of attention heads for each attention layer in the

53 - Transformer encoder.

54 - n_inner (`int`, *optional*, defaults to None):

55 - Dimensionality of the inner feed-forward layers. `None` will set

56 - it to 4 times n_embd

57 - activation_function (`str`, *optional*, defaults to `"gelu"`):

58 - Activation function, to be selected in the list

59 - `["relu", "silu", "gelu", "tanh", "gelu_new", "swiglu"]`.

60 - resid_pdrop (`float`, *optional*, defaults to 0.1):

61 - The dropout probability for all fully connected layers in

62 - the embeddings, encoder, and pooler.

63 - embd_pdrop (`float`, *optional*, defaults to 0.1):

64 - The dropout ratio for the embeddings.

65 - attn_pdrop (`float`, *optional*, defaults to 0.1):

66 - The dropout ratio for the attention.

67 - layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):

68 - The epsilon to use in the layer normalization layers.

69 - initializer_range (`float`, *optional*, defaults to 0.02):

70 - The standard deviation of the truncated_normal_initializer for

71 - initializing all weight matrices.

72 - scale_attn_weights (`bool`, *optional*, defaults to `True`):

73 - Scale attention weights by dividing by sqrt(hidden_size)..

74 - use_cache (`bool`, *optional*, defaults to `True`):

75 - Whether or not the model should return the last key/values

76 - attentions (not used by all models).

77 - scale_attn_by_inverse_layer_idx (`bool`, *optional*, default `True`):

78 - Whether to additionally scale attention weights

79 - by `1 / layer_idx + 1`.

80 - reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`):

81 - Whether to scale keys (K) prior to computing attention

82 - (dot-product)

83 - and upcast attention dot-product/softmax to float() when training

84 - with mixed precision.

85 - position_embedding_type (`str`, *optional*, defaults to `"learned"`):

86 - Positional embedding can be either `"alibi"` or `"learned"`.

87 - mup_width_scale (`float`, *optional*, defaults to 1.0):

88 - muP parameter to scale learning rate and initializers. Calculated

89 - as (`d_model,0 / d_model`), where

90 - `d_model` is the model's width and `d_model,0` is the proxy

91 - model's width.

92 - mup_embeddings_scale (`float`, *optional*, defaults to 1.0):

93 - muP parameter to scale token and position embeddings.

94 - mup_output_alpha (`float`, *optional*, defaults to 1.0):

95 - muP parameter to scale output logits

96 - (`output_logits_scale = mup_output_alpha * mup_width_scale`).

97 - mup_scale_qk_dot_by_d (`bool`, *optional*, defaults to `False`):

98 - Scale attention weights by dividing by hidden_size instead of

99 - sqrt(hidden_size). Need to set scale_attn_weights to `True` as

100 - well.

101 - alibi_scaling (`dict`, *optional*):

102 - Dictionary containing the scaling configuration for ALiBi

103 - embeddings. Currently only supports linear

104 - scaling strategy. Can specify either the scaling `factor` (must be

105 - a float greater than 1) for fixed scaling

106 - or `train_seq_len` for dynamic scaling on input samples with

107 - sequence length > `train_seq_len`. The expected

108 - formats are `{"type": strategy name, "factor": scaling factor}` or

109 - `{"type": strategy name,

110 - "train_seq_len": training sequence length}`.

111 - architectures (`list`, *optional*, defaults to ['JAISLMHeadModel']):

112 - architecture names for Jais.

113 -

114 - Example:

115 -

116 - ```python

117 - >>> from transformers import JAISConfig, JAISModel

118 -

119 - >>> # Initializing a JAIS configuration

120 - >>> configuration = JAISConfig()

121 -

122 - >>> # Initializing a model (with random weights) from the configuration

123 - >>> model = JAISModel(configuration)

124 -

125 - >>> # Accessing the model configuration

126 - >>> configuration = model.config

127 - ```"""

128 -

129 - model_type = "jais"

130 - keys_to_ignore_at_inference = ["past_key_values"]

131 - attribute_map = {

132 - "hidden_size": "n_embd",

133 - "max_position_embeddings": "n_positions",

134 - "num_attention_heads": "n_head",

135 - "num_hidden_layers": "n_layer",

136 - }

137 -

138 - def __init__(

139 - self,

140 - vocab_size=50257,

141 - n_positions=1024,

142 - n_embd=768,

143 - n_layer=12,

144 - n_head=12,

145 - n_inner=None,

146 - activation_function="gelu_new",

147 - resid_pdrop=0.1,

148 - embd_pdrop=0.1,

149 - attn_pdrop=0.1,

150 - layer_norm_epsilon=1e-5,

151 - initializer_range=0.02,

152 - scale_attn_weights=True,

153 - use_cache=True,

154 - bos_token_id=50256,

155 - eos_token_id=50256,

156 - scale_attn_by_inverse_layer_idx=False,

157 - reorder_and_upcast_attn=False,

158 - position_embedding_type="learned",

159 - mup_width_scale=1.0,

160 - mup_embeddings_scale=1.0,

161 - mup_output_alpha=1.0,

162 - mup_scale_qk_dot_by_d=False,

163 - alibi_scaling=None,

164 - architectures=None,

165 - **kwargs,

166 - ):

167 - self.vocab_size = vocab_size

168 - self.n_positions = n_positions

169 - self.n_embd = n_embd

170 - self.n_layer = n_layer

171 - self.n_head = n_head

172 - self.n_inner = n_inner

173 - self.activation_function = activation_function

174 - self.resid_pdrop = resid_pdrop

175 - self.embd_pdrop = embd_pdrop

176 - self.attn_pdrop = attn_pdrop

177 - self.layer_norm_epsilon = layer_norm_epsilon

178 - self.initializer_range = initializer_range

179 - self.scale_attn_weights = scale_attn_weights

180 - self.use_cache = use_cache

181 - self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx

182 - self.reorder_and_upcast_attn = reorder_and_upcast_attn

183 -

184 - self.bos_token_id = bos_token_id

185 - self.eos_token_id = eos_token_id

186 -

187 - self.position_embedding_type = position_embedding_type

188 - self.mup_width_scale = mup_width_scale

189 - self.mup_embeddings_scale = mup_embeddings_scale

190 - self.mup_output_alpha = mup_output_alpha

191 - self.mup_scale_qk_dot_by_d = mup_scale_qk_dot_by_d

192 -

193 - self.alibi_scaling = alibi_scaling

194 - self._alibi_scaling_validation()

195 - if architectures is None:

196 - architectures = ["JAISLMHeadModel"]

197 -

198 - super().__init__(

199 - bos_token_id=bos_token_id,

200 - eos_token_id=eos_token_id,

201 - architectures=architectures,

202 - **kwargs,

203 - )

204 -

205 - def _alibi_scaling_validation(self):

206 - """

207 - Validate the `alibi_scaling` configuration.

208 - """

209 - if self.alibi_scaling is None:

210 - return

211 -

212 - if not isinstance(self.alibi_scaling, dict) or len(self.alibi_scaling) != 2:

213 - raise ValueError(

214 - "`alibi_scaling` must be a dictionary with two fields, "

215 - "`type` and `factor` or `type` and `train_seq_len`, "

216 - f"got {self.alibi_scaling}"

217 - )

218 - alibi_scaling_type = self.alibi_scaling.get("type", None)

219 - alibi_scaling_factor = self.alibi_scaling.get("factor", None)

220 - alibi_dynamic_scaling = self.alibi_scaling.get("train_seq_len", None)

221 - if alibi_scaling_type is None or alibi_scaling_type != "linear":

222 - raise ValueError(

223 - f"`alibi_scaling`'s type field must be 'linear', "

224 - f"got {alibi_scaling_type}"

225 - )

226 - if (

227 - alibi_scaling_factor is not None

228 - and not isinstance(alibi_scaling_factor, float)

229 - or (alibi_scaling_factor is not None and alibi_scaling_factor <= 1.0)

230 - ):

231 - raise ValueError(

232 - f"`alibi_scaling`'s factor field must be a float > 1.0, "

233 - f"got {alibi_scaling_factor}"

234 - )

235 - if (

236 - alibi_dynamic_scaling is not None

237 - and not isinstance(alibi_dynamic_scaling, int)

238 - or (alibi_dynamic_scaling is not None and alibi_dynamic_scaling <= 1)

239 - ):

240 - raise ValueError(

241 - f"`alibi_scaling`'s `train_seq_len` field must be an "

242 - f"integer > 1, got {alibi_dynamic_scaling}"

243 - )