Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions src/transformers/generation/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2140,6 +2140,25 @@ def _tensor_or_none(token, device=None):
generation_config._pad_token_tensor = pad_token_tensor
generation_config._decoder_start_token_tensor = decoder_start_token_tensor


def _is_dynamo_compilation_disabled(self):
"""
Check if dynamo compilation should be disabled based on environment variables.
Uses standard environment variables to respect user's explicit intention to disable
torch.dynamo compilation.
"""
disable_env_vars = [
"TORCHDYNAMO_DISABLE" # Pytorch standard
]

for env_var in disable_env_vars:
value = os.getenv(env_var, "").lower()
if value in ("1", "true", "yes", "on"):
return True

return False


def _valid_auto_compile_criteria(self, model_kwargs: dict[str, Any], generation_config: GenerationConfig) -> bool:
"""
Determines whether to trigger auto-compilation of the model's forward pass at generation time.
Expand All @@ -2148,6 +2167,10 @@ def _valid_auto_compile_criteria(self, model_kwargs: dict[str, Any], generation_
if generation_config.disable_compile:
return False

# honor accelerate config to disable dynamo
if self._is_dynamo_compilation_disabled():
return False

# Base logic
valid_hardware = self.device.type in ["cuda", "xpu"] or bool(
generation_config.compile_config is not None and generation_config.compile_config._compile_all_devices
Expand Down