Skip to content

Commit bf91e1a

Browse files
authored
[Transform] SpinQuant R4 (#1746)
## Purpose ## * Support R4 transforms before R3. R3 requires hooking into the attention module, where as R4 does not ## Prerequisites ## * vllm-project/vllm#22486 ## Testing ## * Performed sanity checks with HF and vLLM --------- Signed-off-by: Kyle Sayers <[email protected]>
1 parent 6af0778 commit bf91e1a

File tree

1 file changed

+18
-4
lines changed
  • src/llmcompressor/modifiers/transform/spinquant

1 file changed

+18
-4
lines changed

src/llmcompressor/modifiers/transform/spinquant/base.py

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ class SpinQuantModifier(Modifier, use_enum_values=True):
104104
@field_validator("randomize", "learnable", mode="before")
105105
def validate_not_implemented(cls, value, info: ValidationInfo):
106106
if value:
107-
raise NotImplementedError(f"{info.field_name} is not supported right now")
107+
raise NotImplementedError(f"{info.field_name} is not supported as of now")
108108
return value
109109

110110
@field_validator("rotations", mode="before")
@@ -237,10 +237,24 @@ def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
237237

238238
def _create_r3_scheme(self) -> TransformScheme:
239239
raise NotImplementedError(
240-
"SpinQuant R3 and R4 rotations will be added in a future release"
240+
"SpinQuant R3 rotations will be added in a future release"
241241
)
242242

243243
def _create_r4_scheme(self) -> TransformScheme:
244-
raise NotImplementedError(
245-
"SpinQuant R3 and R4 rotations will be added in a future release"
244+
return TransformScheme(
245+
type=self.transform_type,
246+
randomize=self.randomize,
247+
requires_grad=self.learnable,
248+
precision=self.precision,
249+
apply=[
250+
TransformArgs(
251+
targets=[*self.mappings.mlp_out],
252+
location="input",
253+
),
254+
TransformArgs(
255+
targets=[*self.mappings.mlp_out],
256+
location="weight_input",
257+
inverse=True,
258+
),
259+
],
246260
)

0 commit comments

Comments
 (0)