File size: 3,974 Bytes
b5dfaf4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
# coding=utf-8
# Copyright 2024 Nvidia Corporation. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import dataclasses
import warnings
from dataclasses import dataclass, MISSING
from functools import partial
from typing import Optional, Dict, Any
from .transformers_4_44_2__configuration_llama import LlamaConfig
from .transformers_4_44_2__modeling_rope_utils import \
rope_config_validation # fake import to make AutoConfig infer the dependency
class DeciLMConfig(LlamaConfig):
model_type = "nemotron-nas"
def __init__(
self,
block_configs: list[dict] | list["BlockConfig"] = None,
**kwargs,
):
super().__init__(**kwargs)
self.intermediate_size = None
self.num_key_value_heads = None
if block_configs is not None:
assert len(block_configs) == self.num_hidden_layers
if isinstance(block_configs[0], dict):
block_configs = [BlockConfig(**conf) for conf in block_configs]
self.block_configs: list[BlockConfig] = block_configs
def to_dict(self) -> Dict[str, Any]:
self_dict = super().to_dict()
if self.block_configs is not None:
self_dict["block_configs"] = [dataclasses.asdict(conf) for conf in self.block_configs]
return self_dict
@partial(dataclass, frozen=True, eq=True, unsafe_hash=True, order=True)
class AttentionConfig:
no_op: bool = False
replace_with_linear: bool = False
n_heads_in_group: Optional[int] = None
def __post_init__(self):
assert not (self.no_op and self.replace_with_linear)
if self.no_op or self.replace_with_linear:
object.__setattr__(self, 'n_heads_in_group', None) # __setattr__ to overcome frozen=True
else:
assert self.n_heads_in_group is not None
@partial(dataclass, frozen=True, eq=True, unsafe_hash=True, order=True)
class FFNConfig:
no_op: bool = False
replace_with_linear: bool = False
ffn_mult: Optional[float] = None
def __post_init__(self):
assert not (self.no_op and self.replace_with_linear)
if self.no_op or self.replace_with_linear:
object.__setattr__(self, 'ffn_mult', None) # __setattr__ to overcome frozen=True
else:
assert self.ffn_mult is not None
@partial(dataclass, frozen=True, eq=True, unsafe_hash=True, order=True)
class BlockConfig:
attention: AttentionConfig = MISSING
ffn: FFNConfig = MISSING
def __post_init__(self):
"""
Init subblock dataclasses from dicts
"""
for subblock_name in dataclasses.fields(self):
subblock_config = getattr(self, subblock_name.name)
if isinstance(subblock_config, dict):
subblock_fields = [field.name for field in dataclasses.fields(subblock_name.type)]
unsupported_fields = [field_name for field_name in subblock_config.keys()
if field_name not in subblock_fields]
if len(unsupported_fields) > 0:
warnings.warn(f"Removed unsupported fields {unsupported_fields} from {subblock_name.type.__name__}")
subblock_config = {k: v for k, v in subblock_config.items() if k not in unsupported_fields}
object.__setattr__(self, subblock_name.name,
subblock_name.type(**subblock_config)) # __setattr__ to overcome frozen=True
|