Fix missing embed_tokens (#671)

casper-hansen · web-flow · commit 8c5dcc6e3d89 · 2024-12-06T18:55:58.000+01:00
diff --git a/awq/models/base.py b/awq/models/base.py
@@ -309,6 +309,7 @@ def forward(self, x):
             max_shard_size=shard_size,
             safe_serialization=safetensors,
             force_contiguous=True,
+            shared_tensors_to_discard=self.model._tied_weights_keys,
         )
 
     @classmethod
diff --git a/setup.py b/setup.py
@@ -42,6 +42,7 @@
     "accelerate",
     "datasets>=2.20",
     "zstandard",
+    "huggingface_hub @ git+https://github.com/huggingface/huggingface_hub@fix-discard-shared-tensors",
 ]
 
 setup(

Original file line number	Diff line number	Diff line change
`@@ -309,6 +309,7 @@ def forward(self, x):`
`309`	`309`	`max_shard_size=shard_size,`
`310`	`310`	`safe_serialization=safetensors,`
`311`	`311`	`force_contiguous=True,`
	`312`	`+ shared_tensors_to_discard=self.model._tied_weights_keys,`
`312`	`313`	`)`
`313`	`314`
`314`	`315`	`@classmethod`
Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@`
`42`	`42`	`"accelerate",`
`43`	`43`	`"datasets>=2.20",`
`44`	`44`	`"zstandard",`
	`45`	`+ "huggingface_hub @ git+https://github.com/huggingface/huggingface_hub@fix-discard-shared-tensors",`
`45`	`46`	`]`
`46`	`47`
`47`	`48`	`setup(`