UNSLOTH_DISABLE_STATIC_GENERATION

This commit is contained in:
Daniel Han 2025-05-15 04:48:42 -07:00
commit 2fe9067d6b
2 changed files with 5 additions and 4 deletions

View file

@ -542,8 +542,7 @@ class FastModel(FastBaseModel):
if transformers_version < Version("4.50.0.dev0"):
raise RuntimeError("Unsloth: Granite Vision only works on transformers >= 4.50.0." + NIGHTLY)
elif "csm-1b" in model_name.lower():
os.environ["UNSLOTH_COMPILE_DISABLE"] = "0"
os.environ["UNSLOTH_DISABLE_FAST_GENERATION"] = "0"
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "torch.float16;if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
elif "olmo-2" in model_name.lower() and transformers_version < Version("4.50.0.dev0"):
raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY)

View file

@ -188,7 +188,10 @@ def unsloth_base_fast_generate(
# Use hybrid if sliding window seen, otherwise try static
cache_implementation = getattr(self.config, "cache_implementation", None)
if getattr(self, "_supports_static_cache", True):
cache_implementation = "static"
if os.environ.get("UNSLOTH_DISABLE_STATIC_GENERATION", "0") == "0":
cache_implementation = "static"
else:
cache_implementation = None
else:
cache_implementation = None
if cache_implementation is not None:
@ -204,7 +207,6 @@ def unsloth_base_fast_generate(
kwargs["cache_implementation"] = cache_implementation
kwargs["compile_config"] = _compile_config
pass
del kwargs["cache_implementation"]
try:
with torch.inference_mode(), autocaster: