Update dynamic env handling to preserve None when USE_DYNAMIC is unset

[yafshar]

What does this PR do?

Modified the logic for setting self.dynamic to explicitly preserve None when the USE_DYNAMIC environment variable is not set, aligning with the behavior described in the PyTorch documentation for torch.compile (https://docs.pytorch.org/stable/generated/torch.compile.html). The documentation notes that dynamic=None has distinct semantics from dynamic=False, where None indicates a different configuration state. Previously, the code defaulted to False when the environment variable was unset, which could lead to incorrect behavior.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[yafshar]

@IlyasMoutawwakil can you check this PR?

[yafshar]

You can try the change by a simple test and also observe the old behavior as it is

import unittest
import os
from dataclasses import dataclass, field
from unittest.mock import patch
import copy

def str_to_bool(value):
    if value.lower() in ("y", "yes", "t", "true", "on", "1"):
        return 1

class KwargsHandler:
    def to_dict(self):
        return copy.deepcopy(self.__dict__)

    def to_kwargs(self):
        default_dict = self.__class__().to_dict()
        this_dict = self.to_dict()
        return {k: v for k, v in this_dict.items() if default_dict[k] != v}

@dataclass
class TorchDynamoPluginCurrent(KwargsHandler):
    dynamic: bool = field(default=None, metadata={"help": "Whether to use dynamic shape"})

    def __post_init__(self):
        prefix = "ACCELERATE_DYNAMO_"
        if self.dynamic is None:
            self.dynamic = str_to_bool(os.environ.get(prefix + "USE_DYNAMIC", "False")) == 1

    def to_dict(self):
        dynamo_config = copy.deepcopy(self.__dict__)
        return dynamo_config

    def to_kwargs(self):
        kwargs = super().to_kwargs()
        return kwargs

@dataclass
class TorchDynamoPluginModified(KwargsHandler):
    dynamic: bool = field(default=None, metadata={"help": "Whether to use dynamic shape"})

    def __post_init__(self):
        prefix = "ACCELERATE_DYNAMO_"
        if self.dynamic is None:
            env_value = os.environ.get(prefix + "USE_DYNAMIC")
            self.dynamic = (str_to_bool(env_value) == 1) if env_value is not None else None

    def to_dict(self):
        dynamo_config = copy.deepcopy(self.__dict__)
        return dynamo_config

    def to_kwargs(self):
        kwargs = super().to_kwargs()
        return kwargs

class TestTorchDynamoPlugin(unittest.TestCase):
    def test_dynamic_current_implementation(self):
        """Test the current implementation of TorchDynamoPlugin.__post_init__"""
        # Test case 1: Environment variable set to True
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "True"}):
            plugin = TorchDynamoPluginCurrent()
            self.assertTrue(plugin.dynamic, "Expected dynamic to be True when env is 'True'")

        # Test case 2: Environment variable set to False
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "False"}):
            plugin = TorchDynamoPluginCurrent()
            self.assertFalse(plugin.dynamic, "Expected dynamic to be False when env is 'False'")

        # Test case 3: Environment variable unset (should be None, but current fails)
        with patch.dict(os.environ, {}, clear=True):
            plugin = TorchDynamoPluginCurrent()
            self.assertFalse(plugin.dynamic, "Current implementation incorrectly sets dynamic to False when env is unset")

    def test_dynamic_modified_implementation(self):
        """Test the modified implementation of TorchDynamoPlugin.__post_init__"""
        # Test case 1: Environment variable set to True
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "True"}):
            plugin = TorchDynamoPluginModified()
            self.assertTrue(plugin.dynamic, "Expected dynamic to be True when env is 'True'")

        # Test case 2: Environment variable set to False
        with patch.dict(os.environ, {"ACCELERATE_DYNAMO_USE_DYNAMIC": "False"}):
            plugin = TorchDynamoPluginModified()
            self.assertFalse(plugin.dynamic, "Expected dynamic to be False when env is 'False'")

        # Test case 3: Environment variable unset (should be None)
        with patch.dict(os.environ, {}, clear=True):
            plugin = TorchDynamoPluginModified()
            self.assertIsNone(plugin.dynamic, "Expected dynamic to be None when env is unset")

if __name__ == "__main__":
    unittest.main()

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Indeed thanks for spotting this. Evenutually, we also need to fix for USE_DYNAMIC also to align the behavior