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

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 14, 05:30 UTC

Support RL online quantization with torchao (#23014)

Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>

Jerry Zhang committed Oct 1, 2025 at 16:39 UTC c31246800c98272ab28b7d454d6b72f38e396972

6 files changed +465 -16

tests/quantization/test_torchao.py

+121 -5

@@ -20,7 +20,6 @@ def test_pre_quantized_model(vllm_runner):

20 output = llm.generate_greedy(["The capital of France is"],

21 max_tokens=32)

22 assert output

23 - print(output)

25 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

@@ -42,7 +41,6 @@ def test_opt_125m_int8wo_model_loading_with_params(vllm_runner,

41 max_tokens=32)

43 assert output

45 - print(output)

46 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

@@ -57,7 +55,6 @@ def test_opt_125m_int4wo_model_per_module_quant(vllm_runner):

55 max_tokens=32)

57 assert output

60 - print(output)

60 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

@@ -72,7 +69,6 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):

69 max_tokens=32)

71 assert output

75 - print(output)

74 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

@@ -92,7 +88,127 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):

88 max_tokens=32)

90 assert output

95 - print(output)

91 +

92 +

93 +@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

94 +def test_on_the_fly_quant_config_dict_json(vllm_runner):

95 + """Testing on the fly quantization, load_weights integration point,

96 + with config dict serialized to json string

97 + """

98 + torch._dynamo.reset()

99 + model_name = "facebook/opt-125m"

100 +

101 + import json

102 +

103 + from torchao.core.config import config_to_dict

104 + from torchao.quantization import (

105 + Float8DynamicActivationFloat8WeightConfig, PerRow)

106 +

107 + torchao_quant_config = Float8DynamicActivationFloat8WeightConfig(

108 + granularity=PerRow())

109 + hf_overrides = {

110 + "quantization_config_dict_json":

111 + json.dumps(config_to_dict(torchao_quant_config))

112 + }

113 + with vllm_runner(model_name=model_name,

114 + dtype="bfloat16",

115 + pt_load_map_location="cuda:0",

116 + quantization="torchao",

117 + hf_overrides=hf_overrides) as llm:

118 + output = llm.generate_greedy(["The capital of France is"],

119 + max_tokens=32)

120 +

121 + assert output

122 +

123 +

124 +@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

125 +def test_on_the_fly_quant_config_file(vllm_runner):

126 + """Testing on the fly quantization, load_weights integration point,

127 + with config file

128 + """

129 + torch._dynamo.reset()

130 + model_name = "facebook/opt-125m"

131 + import json

132 + from tempfile import NamedTemporaryFile

133 +

134 + from torchao.core.config import config_to_dict

135 + from torchao.quantization import (

136 + Float8DynamicActivationFloat8WeightConfig, PerRow)

137 +

138 + config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())

139 +

140 + with NamedTemporaryFile(mode="w", delete=False) as f:

141 + f.write(json.dumps(config_to_dict(config)))

142 + # close the file to save it

143 + f.close()

144 + config_file_name = str(f.name)

145 +

146 + hf_overrides = {"quantization_config_file": config_file_name}

147 + with vllm_runner(model_name=model_name,

148 + dtype="bfloat16",

149 + pt_load_map_location="cuda:0",

150 + quantization="torchao",

151 + hf_overrides=hf_overrides) as llm:

152 + output = llm.generate_greedy(["The capital of France is"],

153 + max_tokens=32)

154 +

155 + assert output

156 +

157 +

158 +@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

159 +def test_reload_weights():

160 + import json

161 +

162 + from torchao.core.config import config_to_dict

163 + from torchao.quantization import (

164 + Float8DynamicActivationFloat8WeightConfig, PerRow)

165 +

166 + from vllm import LLM, SamplingParams

167 +

168 + torchao_quant_config = Float8DynamicActivationFloat8WeightConfig(

169 + granularity=PerRow())

170 +

171 + hf_overrides = {

172 + "quantization_config_dict_json":

173 + json.dumps(config_to_dict(torchao_quant_config))

174 + }

175 +

176 + llm = LLM(

177 + model="Qwen/Qwen3-0.6B",

178 + dtype="bfloat16",

179 + load_format="dummy",

180 + enforce_eager=True,

181 + quantization="torchao",

182 + hf_overrides=hf_overrides,

183 + )

184 + # Update load format from `dummy` to `auto`

185 + llm.collective_rpc("update_config",

186 + args=({

187 + "load_config": {

188 + "load_format": "auto"

189 + }

190 + }, ))

191 + # Now reload real weights inplace

192 + llm.collective_rpc("reload_weights")

193 + prompts = [

194 + "Hello, my name is",

195 + "The president of the United States is",

196 + "The capital of France is",

197 + "The future of AI is",

198 + ]

199 + # Create a sampling params object.

200 + sampling_params = SamplingParams(temperature=0, top_p=0.95)

201 + outputs = llm.generate(prompts, sampling_params)

202 + # make sure it runs

203 + for output in outputs:

204 + generated_text = output.outputs[0].text

205 + assert generated_text

206 + # can also uncomment locally to make sure the generated

207 + # output makes sense

208 + # prompt = output.prompt

209 + # print(f"Prompt: {prompt!r}")

210 + # print(f"Output: {generated_text!r}")

211 + # print("-" * 60)

212

213

214 if __name__ == "__main__":

vllm/model_executor/layers/quantization/torchao.py

+64 -8

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +import json

4 from typing import Any, Optional

6 import torch

@@ -40,7 +41,8 @@ class TorchAOConfig(QuantizationConfig):

42 def __init__(self,

43 torchao_config,

43 - skip_modules: Optional[list[str]] = None) -> None:

44 + skip_modules: Optional[list[str]] = None,

45 + is_checkpoint_torchao_serialized: bool = False) -> None:

46 """

47 # TorchAO quantization relies on tensor subclasses. In order,

48 # to enable proper caching this needs standalone compile

@@ -58,9 +60,11 @@ class TorchAOConfig(QuantizationConfig):

60 super().__init__()

61 self.torchao_config = torchao_config

62 self.skip_modules = skip_modules or []

63 + self.is_checkpoint_torchao_serialized = is_checkpoint_torchao_serialized

65 def __repr__(self) -> str:

63 - return f"TorchAOConfig({self.torchao_config})"

66 + return f"TorchAOConfig({self.torchao_config=}, {self.skip_modules=}, " \

67 + f"{self.is_checkpoint_torchao_serialized=})"

69 def get_name(self) -> QuantizationMethods:

70 return "torchao"

@@ -74,7 +78,10 @@ class TorchAOConfig(QuantizationConfig):

79 @staticmethod

80 def get_config_filenames() -> list[str]:

77 - return ["config.json"]

81 + """torchao doesn't require additional config files, we use

82 + `config.json` from huggingface: `model_config.hf_config`

83 + """

84 + return []

86 @classmethod

87 def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig":

@@ -87,6 +94,10 @@ class TorchAOConfig(QuantizationConfig):

94 "`pip install torchao>=0.10.0` to use torchao quantization."

95 ) from err

97 + quant_method = cls.get_from_keys_or(config, ["quant_method"], None)

98 + is_checkpoint_torchao_serialized = (quant_method is not None

99 + and "torchao" in quant_method)

100 +

101 hf_config = cls.get_from_keys_or(config, ["quant_type"], None)

102 assert hf_config is not None, "quant_type must be specified"

103 assert len(hf_config) == 1 and "default" in hf_config, (

@@ -110,7 +121,38 @@ class TorchAOConfig(QuantizationConfig):

121 if layer_cfg is None:

122 skip_modules.append(layer)

123

113 - return cls(ao_config, skip_modules)

124 + return cls(ao_config, skip_modules, is_checkpoint_torchao_serialized)

125 +

126 + @classmethod

127 + def from_config_file(cls, config_file: str) -> "TorchAOConfig":

128 + """Initialize class from a config file. Example:

129 + ```

130 + config = (

131 + Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())

132 + )

133 + fn = "torchao_config.json"

134 +

135 + with open(fn, "w") as f:

136 + f.write(json.dumps(config_to_dict(config)))

137 + ```

138 + """

139 + with open(config_file) as f:

140 + f.seek(0)

141 + f_read = f.read()

142 + config_dict = json.loads(f_read)

143 +

144 + hf_config = {"quant_type": {"default": config_dict}}

145 + return cls.from_config(hf_config)

146 +

147 + @classmethod

148 + def from_config_dict_json(cls, config_dict_json: str) -> "TorchAOConfig":

149 + """Iniitalize class from a config_dict json string, got from

150 + torchao_config_object = some AOBaseConfig object

151 + json.dumps(config_to_dict(torchao_config_object))

152 + """

153 + config_dict = json.loads(config_dict_json)

154 + hf_config = {"quant_type": {"default": config_dict}}

155 + return cls.from_config(hf_config)

156

157 def get_quant_method(self, layer: torch.nn.Module,

158 prefix: str) -> Optional["QuantizeMethodBase"]:

@@ -128,7 +170,9 @@ class TorchAOConfig(QuantizationConfig):

170 c = module_fqn_to_config.get(

171 module_fqn) or module_fqn_to_config.get("_default", None)

172 if c is not None:

131 - current_torchao_config = TorchAOConfig(c, self.skip_modules)

173 + current_torchao_config = TorchAOConfig(

174 + c, self.skip_modules,

175 + self.is_checkpoint_torchao_serialized)

176 return TorchAOLinearMethod(current_torchao_config)

177 else:

178 return UnquantizedLinearMethod()

@@ -172,7 +216,7 @@ class TorchAOLinearMethod(LinearMethodBase):

216 """Linear method for torchao.

217

218 Args:

175 - quant_config: The torchao quantization config, a string that encodes

219 + quant_config: The torchao quantization config, a string that encodes

220 the type of quantization and all relevant arguments.

221 """

222

@@ -197,8 +241,9 @@ class TorchAOLinearMethod(LinearMethodBase):

241 ),

242 requires_grad=False,

243 )

200 - weight = torchao_quantize_param_data(weight,

201 - self.quant_config.torchao_config)

244 + if self.quant_config.is_checkpoint_torchao_serialized:

245 + weight = torchao_quantize_param_data(

246 + weight, self.quant_config.torchao_config)

247

248 set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})

249

@@ -212,3 +257,14 @@ class TorchAOLinearMethod(LinearMethodBase):

257 bias: Optional[torch.Tensor] = None,

258 ) -> torch.Tensor:

259 return F.linear(x, layer.weight, bias)

260 +

261 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

262 + if self.quant_config.is_checkpoint_torchao_serialized:

263 + return

264 +

265 + # quantize the weight on the fly if the checkpoint is not already

266 + # quantized by torchao

267 + weight = torchao_quantize_param_data(layer.weight,

268 + self.quant_config.torchao_config)

269 + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})

270 + layer.register_parameter("weight", weight)

vllm/model_executor/model_loader/default_loader.py

+29 -2

@@ -261,8 +261,35 @@ class DefaultModelLoader(BaseModelLoader):

261 def load_weights(self, model: nn.Module,

262 model_config: ModelConfig) -> None:

263 weights_to_load = {name for name, _ in model.named_parameters()}

264 - loaded_weights = model.load_weights(

265 - self.get_all_weights(model_config, model))

264 +

265 + # if we don't have `model.weight_metadata_and_attr_saved` defined and

266 + # set to True, it means that this is either offline quantization case

267 + # or the first run of online quantization

268 + # see online_quantization.py for detailed notes

269 + offline_quantization_or_first_run_of_online_quantization = not getattr(

270 + model, "weight_metadata_and_attr_saved", False)

271 +

272 + if model_config.quantization is None:

273 + # model is not quantized

274 + loaded_weights = model.load_weights(

275 + self.get_all_weights(model_config, model))

276 + elif offline_quantization_or_first_run_of_online_quantization:

277 + # case 1: offline quantized checkpoint

278 + # case 2: Step I1 first run of weight loading with

279 + # online quantization

280 + # see online_quantization.py for detailed notes

281 + loaded_weights = model.load_weights(

282 + self.get_all_weights(model_config, model))

283 + else:

284 + # to avoid circular dependency

285 + from vllm.model_executor.model_loader.online_quantization import (

286 + load_weights_and_online_quantize)

287 +

288 + # subsequent runs of weight loading with online

289 + # quantization

290 + loaded_weights = load_weights_and_online_quantize(

291 + self, model, model_config)

292 +

293 self.counter_after_loading_weights = time.perf_counter()

294 logger.info(

295 "Loading weights took %.2f seconds",

vllm/model_executor/model_loader/online_quantization.py new

+217

@@ -0,0 +1,217 @@

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

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

3 +

4 +import types

5 +

6 +import torch

7 +from torch import nn

8 +

9 +from vllm.config import ModelConfig

10 +from vllm.logger import init_logger

11 +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader

12 +from vllm.model_executor.model_loader.utils import (

13 + process_weights_after_loading)

14 +

15 +logger = init_logger(__name__)

16 +

17 +# Notes for Online Quantization

18 +# In terms of state of checkpoints, quantization config and their

19 +# correspondance to online quantization:

20 +# | Use Case | Checkpoints | model_config.quantization |

21 +# | no quant | high precision | None |

22 +# | offline quant | quantized | fp8, torchao etc. |

23 +# | online quant | high precision | torchao etc. |

24 +#

25 +# The process for loading non-quantized checkpoint

26 +# 1. load non-quantized weights (load_weights)

27 +# 2. do any additional post processing (process_weights_after_loading)

28 +#

29 +# The process for loading offline quantized checkpoint

30 +# 1. load offline-quantized weights (load_weights)

31 +# 2. do any additional post processing (process_weights_after_loading)

32 +

33 +# The process for unquantized model reloading

34 +# (repeated run in RL training loop)

35 +# first run

36 +# UI1. load_weights: load bfloat16 weights

37 +# UI2. process_weights_after_loading: any additional post processing

38 +# subsequent run

39 +# UC1: load_weights: load bfloat16 weights

40 +# (shouldn't be any issues since we didn't change any attributes

41 +# of the weights)

42 +# UC2: process_weights_after_loading: any additional post processing

43 +

44 +# The process for weight reloading with online quantization

45 +# (repeated run in RL training loop)

46 +# first run

47 +# I1. load_weights: load bfloat16 weights

48 +# I2. process_weights_after_loading:

49 +# record weight metadata and attributes for R1 and R2

50 +# quantize weights to fp8

51 +# subsequent run

52 +# (beginning model weight is in fp8)

53 +# load_weights:

54 +# R1. restore bfloat16 model weight metadata

55 +# R2. restore the model weight attributes

56 +# R3. reload bfloat16 weights

57 +# R4. quantize weights (by calling process_weights_after_loading),

58 +# also set `process_weights_after_loading_already_called` to

59 +# True to stop it from running again

60 +# process_weights_after_loading (if called):

61 +# this will be skipped since it's already ran in

62 +# load_weights

63 +

64 +

65 +def maybe_save_metadata_and_attributes_for_weight_reloading(

66 + model: nn.Module, model_config: ModelConfig):

67 + # following is to support on the fly quantization, currently only supported

68 + # for torchao

69 + if model_config.quantization != "torchao":

70 + return

71 +

72 + if getattr(model, "process_weights_after_loading_already_called", False):

73 + # In case `process_weights_after_loading` is called multiple times

74 + # we'll skip it at later times

75 + logger.warning(

76 + "process_weights_after_loading already called for model %s", model)

77 + return

78 +

79 + from vllm.model_executor.model_loader.weight_utils import get_quant_config

80 + quant_config = get_quant_config(model_config, None)

81 +

82 + # If checkpoint is already torchao serialized, this means it's

83 + # pre-quantized quantization case, we'll skip saving the metadata

84 + # Otherwise, this is Step I2 of initialization steps of

85 + # online quantization

86 + # This step record the weights metadata and weight attributes so we can

87 + # restore the bfloat16 model weights during the relad step (R1 and R2)

88 + # see Notes in online_quantization.py for more details

89 + if not (hasattr(quant_config, "is_checkpoint_torchao_serialized") and \

90 + not quant_config.is_checkpoint_torchao_serialized):

91 + return

92 +

93 + # This is the I2 step of online quantiztion that saves

94 + # metadata and attributes of weights so they can be used in R1 and

95 + # R2 step, note that we only save these during initialization

96 +

97 + # Includes two things

98 + # 1. save floating point metadata (shape, dtype, device) for init

99 + # 2. save weight attributes, e.g. `output_dim`, `weight_loader` for init

100 +

101 + if getattr(model, "weight_metadata_and_attr_saved", False):

102 + return

103 +

104 + # save the dtype, shape and device for model parameter, used for

105 + # restoring the model high precision parameters before

106 + # reloading the weights

107 + assert not hasattr(model, "original_weights_rebuild_keys")

108 + model.original_weights_rebuild_keys = {}

109 + for name, p in model.named_parameters():

110 + model.original_weights_rebuild_keys[name] = {

111 + "shape": p.shape,

112 + "dtype": p.dtype,

113 + "device": p.device,

114 + }

115 +

116 + # record the weight attributes (loader functions etc.)

117 + # so these can be recovered later when we reload the weights

118 + # structure: {"weight_name": {"weight_attr_key": attr}}

119 + assert not hasattr(model, "recorded_weight_attr")

120 + model.recorded_weight_attr = {}

121 + for name, param in model.named_parameters():

122 + model.recorded_weight_attr[name] = {}

123 + for key in param.__dict__:

124 + if hasattr(param, key):

125 + attr = getattr(param, key)

126 + if not callable(attr):

127 + model.recorded_weight_attr[name][key] = attr

128 + elif hasattr(attr, "__self__") and param is attr.__self__:

129 + # if attr is a bonded method for an instance, and

130 + # attr.__self__ points to the instance (param)

131 + # we'll record the underlying function object

132 + model.recorded_weight_attr[name][key] = attr.__func__

133 + else:

134 + model.recorded_weight_attr[name][key] = attr

135 + # mark the metadata and attributes saved so we don't run it again

136 + model.weight_metadata_and_attr_saved = True

137 +

138 +

139 +def _bond_method_to_cls(func, obj):

140 + if hasattr(func, "__self__") or not callable(func):

141 + # If the function is already bound to an instance, return it as is

142 + return func

143 + else:

144 + return types.MethodType(func, obj)

145 +

146 +

147 +def load_weights_and_online_quantize(model_loader: DefaultModelLoader,

148 + model: nn.Module,

149 + model_config: ModelConfig) -> set[str]:

150 + # online quantization, right now only enabled for

151 + # torchao

152 + # R1, R2, R3, R4 in the Notes

153 +

154 + # TODO: Add fp8 support

155 + assert model_config.quantization == "torchao", "online " \

156 + "quantization is only enabled for torchao currently"

157 + # TODO: use create_weights to restore the weights to original state

158 +

159 + # Step R1: First restore the quantized weights to original bfloat16

160 + # weights, with original metadata (shape, dtype, device)

161 + # and attributes, so that bfloat16 weights can be loaded properly

162 + existing_param_names = dict(

163 + model.named_parameters(remove_duplicate=False)).keys()

164 + named_modules = dict(model.named_modules(remove_duplicate=False))

165 + model_device = None

166 +

167 + # Step R2: recover the parameter to the state before first loading

168 + for name, d in model.original_weights_rebuild_keys.items():

169 + _shape = d["shape"]

170 + _dtype = d["dtype"]

171 + _device = d["device"]

172 + if model_device is not None:

173 + assert model_device == _device, "Expecting all weights " \

174 + "to be in the same device for now, got both: " \

175 + f"{model_device} and {_device}"

176 + else:

177 + model_device = _device

178 +

179 + if name in existing_param_names:

180 + module_name, weight_name = name.rsplit(".", 1)

181 + module = named_modules[module_name]

182 + setattr(

183 + module, weight_name,

184 + torch.nn.Parameter(

185 + torch.empty(_shape, dtype=_dtype, device=_device)))

186 +

187 + # recorded_weight_attr is

188 + # {"weight_name": {"weight_attr_key": attr}}

189 + # e.g.

190 + # {

191 + # {

192 + # "layer.0.weight": {

193 + # "weight_loader": weight_loader_function_object,

194 + # "input_dim": 0, ...

195 + # },

196 + # "layer.1.weight": ...,

197 + # }

198 + # }

199 + for full_weight_name, weight_attr_dict in \

200 + model.recorded_weight_attr.items():

201 + for attr_name, attr in weight_attr_dict.items():

202 + module_name, weight_name = full_weight_name.rsplit(".", 1)

203 + module = named_modules[module_name]

204 + weight = getattr(module, weight_name)

205 + if not hasattr(weight, attr_name):

206 + setattr(weight, attr_name, _bond_method_to_cls(attr, weight))

207 +

208 + # Step I1: reload bfloat16 / high precision weights

209 + loaded_weights = model.load_weights(

210 + model_loader.get_all_weights(model_config, model))

211 +

212 + # Step I2: online quantize the weights

213 + # manually process weights after loading

214 + model.process_weights_after_loading_already_called = False

215 + process_weights_after_loading(model, model_config, model_device)

216 + model.process_weights_after_loading_already_called = True

217 + return loaded_weights

vllm/model_executor/model_loader/utils.py

+8 -1

@@ -95,6 +95,13 @@ def initialize_model(

96 def process_weights_after_loading(model: nn.Module, model_config: ModelConfig,

97 target_device: torch.device) -> None:

98 +

99 + # to avoid circular dependency

100 + from vllm.model_executor.model_loader.online_quantization import (

101 + maybe_save_metadata_and_attributes_for_weight_reloading)

102 + maybe_save_metadata_and_attributes_for_weight_reloading(

103 + model, model_config)

104 +

105 for _, module in model.named_modules():

106 if isinstance(module, QKVCrossParallelLinear):

107 # NOTE(Isotr0py): special case for cross QKV layer because

@@ -243,7 +250,7 @@ def get_architecture_class_name(model_config: ModelConfig) -> str:

250 class ParamMapping:

251 """

252 A class to handle parameter mapping for model weight loading.

246 - It creates a bidirectional mapping between packed parameters and their

253 + It creates a bidirectional mapping between packed parameters and their

254 constituent parts.

255 """

256 packed_mapping: dict[str, list[str]]

vllm/model_executor/model_loader/weight_utils.py

+26

@@ -246,8 +246,34 @@ def get_quant_config(model_config: ModelConfig,

246 # compressed-tensors uses a compressions_config

247 hf_quant_config = getattr(model_config.hf_config, "compression_config",

248 None)

249 +

250 if hf_quant_config is not None:

251 return quant_cls.from_config(hf_quant_config)

252 +

253 + # if hf_quant_config is None, we will try to get config from

254 + # hf_overrides

255 + hf_overrides = model_config.hf_overrides

256 + quantization_config_file = hf_overrides.get("quantization_config_file",

257 + None)

258 + if quantization_config_file is not None:

259 + if hasattr(quant_cls, "from_config_file"):

260 + return quant_cls.from_config_file(quantization_config_file)

261 + else:

262 + raise NotImplementedError(

263 + "from_config_file is specified in hf_override config, "

264 + "but quant_cls.from_config_file is not implemented in "

265 + f"{quant_cls}")

266 + quantization_config_json = hf_overrides.get(

267 + "quantization_config_dict_json", None)

268 + if quantization_config_json is not None:

269 + if hasattr(quant_cls, "from_config_dict_json"):

270 + return quant_cls.from_config_dict_json(quantization_config_json)

271 + else:

272 + raise NotImplementedError(

273 + "from_config_dict_json is specified in hf_override config, "

274 + "but quant_cls.from_config_dict_json is not implemented in "

275 + f"{quant_cls}")

276 +

277 # Inflight BNB quantization

278 if model_config.quantization == "bitsandbytes":

279 return quant_cls.from_config({})