Merge branch 'main' into akannan/op_support_maskedscatter

forge/test/mlir/llama/test_llama_backward.py

@@ -2,27 +2,19 @@
 
 # SPDX-License-Identifier: Apache-2.0
 
-import torch
 import pytest
-from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer
 
 import forge
 from test.mlir.llama.utils.utils import load_model
 
 
-# TODO(tt-mlir issue #1503): This test is failing because the embedding op doesn't work with FP32.
-# It should be fixed in the tt-mlir compiler soon.
 @pytest.mark.parametrize("model_path", ["openlm-research/open_llama_3b"])
-@pytest.mark.xfail()
 def test_llama_backward(model_path):
     # Load Model and Tokenizer
     framework_model, tokenizer = load_model(model_path)
 
     prompt = "Q: What is the largest animal?\nA:"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
 
-    loss_fn = torch.nn.CrossEntropyLoss()
-    framework_optimizer = torch.optim.SGD(framework_model.parameters(), lr=1e-3)
-
     # Compile the model with loss and optimizer, this will invoke an autograd pass which produces bwd graph.
-    compiled_model = forge.compile(framework_model, input_ids, loss=loss_fn, optimizer=framework_optimizer)
+    compiled_model = forge.compile(framework_model, input_ids, training=True)

forge/test/operators/utils/compat.py

@@ -13,6 +13,7 @@
 from forge import ForgeModule, Module, DepricatedVerifyConfig
 from forge.tensor import to_pt_tensors
 from forge.op_repo import TensorShape
+from forge.config import CompilerConfig
 from forge.verify.compare import compare_with_golden
 from forge.verify.verify import verify
 from forge.verify.config import VerifyConfig
@@ -303,6 +304,7 @@ def create_torch_inputs(
 def verify_module_for_inputs_deprecated(
     model: Module,
     inputs: List[torch.Tensor],
+    compiler_cfg: CompilerConfig,
     pcc: Optional[float] = None,
     dev_data_format: forge.DataFormat = None,
     convert_to_forge: bool = True,  # explicit conversion to forge data format
@@ -315,7 +317,7 @@ def verify_module_for_inputs_deprecated(
     else:
         forge_inputs = inputs
 
-    compiled_model = forge.compile(model, sample_inputs=forge_inputs)
+    compiled_model = forge.compile(model, sample_inputs=forge_inputs, compiler_cfg=compiler_cfg)
     co_out = compiled_model(*forge_inputs)
 
     # TODO check output data format type
@@ -338,6 +340,7 @@ def verify_module_for_inputs_deprecated(
 def verify_module_for_inputs(
     model: Module,
     inputs: List[torch.Tensor],
+    compiler_cfg: CompilerConfig,
     verify_config: Optional[VerifyConfig] = VerifyConfig(),
     dev_data_format: forge.DataFormat = None,
     convert_to_forge: bool = True,  # explicit conversion to forge data format
@@ -348,7 +351,7 @@ def verify_module_for_inputs(
     else:
         forge_inputs = inputs
 
-    compiled_model = forge.compile(model, sample_inputs=forge_inputs)
+    compiled_model = forge.compile(model, sample_inputs=forge_inputs, compiler_cfg=compiler_cfg)
     verify(inputs, model, compiled_model, verify_config)
 
 

forge/test/operators/utils/utils.py

@@ -22,6 +22,7 @@
 from forge.verify import TestKind  # , verify_module
 from forge._C import MathFidelity
 
+from forge.config import CompilerConfig
 from forge.verify.config import VerifyConfig
 
 from .compat import TestDevice
@@ -93,16 +94,21 @@ class CompilerUtils:
     """Utility functions for Forge compiler configuration"""
 
     @staticmethod
-    def set_input_source(input_source_flag: InputSourceFlag):
+    def set_input_source(input_source_flag: InputSourceFlag, compiler_cfg: CompilerConfig):
         """Set compiler configuration for input source"""
+        # Not existing in the compiler, after global config removal
+        # compiler_cfg.input_queues_on_host = input_source_flag.input_queues_on_host
+        # if input_source_flag.set_default_dram_parameters:
+        #     compiler_cfg.default_dram_parameters = input_source_flag.default_dram_parameters
+
         # NOP since we don't use this flag in the compiler, currently.
         pass
 
     @staticmethod
-    def set_math_fidelity(math_fidelity: MathFidelity):
+    def set_math_fidelity(math_fidelity: MathFidelity, compiler_cfg: CompilerConfig):
         """Set compiler configuration for math fidelity"""
-        # NOP since we don't use this flag in the compiler, currently.
-        pass
+        # Currently not respected/supported in the compiler
+        compiler_cfg.default_math_fidelity = math_fidelity
 
 
 class DeviceUtils:
@@ -155,7 +161,10 @@ def verify(
             skip_forge_verification: Skip verification with Forge module
         """
 
+        compiler_cfg = CompilerConfig()
+
         cls.setup(
+            compiler_cfg=compiler_cfg,
             input_source_flag=input_source_flag,
             math_fidelity=math_fidelity,
             warm_reset=warm_reset,
@@ -172,6 +181,7 @@ def verify(
             cls.verify_module_for_inputs_deprecated(
                 model=model,
                 inputs=inputs,
+                compiler_cfg=compiler_cfg,
                 pcc=pcc,
                 dev_data_format=dev_data_format,
                 convert_to_forge=convert_to_forge,
@@ -186,6 +196,7 @@ def verify(
             cls.verify_module_for_inputs(
                 model=model,
                 inputs=inputs,
+                compiler_cfg=compiler_cfg,
                 verify_config=verify_config,
                 dev_data_format=dev_data_format,
                 convert_to_forge=convert_to_forge,
@@ -194,6 +205,7 @@ def verify(
     @classmethod
     def setup(
         cls,
+        compiler_cfg: CompilerConfig,
         input_source_flag: InputSourceFlags = None,
         math_fidelity: forge.MathFidelity = None,
         warm_reset: bool = False,
@@ -202,10 +214,10 @@ def setup(
             DeviceUtils.warm_reset()
 
         if input_source_flag:
-            CompilerUtils.set_input_source(input_source_flag.value)
+            CompilerUtils.set_input_source(input_source_flag.value, compiler_cfg)
 
         if math_fidelity:
-            CompilerUtils.set_math_fidelity(math_fidelity)
+            CompilerUtils.set_math_fidelity(math_fidelity, compiler_cfg)
 
         # if dev_data_format:
         #     input_params.append({"dev_data_format": dev_data_format})
@@ -233,6 +245,7 @@ def verify_module_for_inputs_deprecated(
         cls,
         model: Module,
         inputs: List[torch.Tensor],
+        compiler_cfg: CompilerConfig,
         pcc: Optional[float] = None,
         dev_data_format: forge.DataFormat = None,
         convert_to_forge: bool = True,  # explicit conversion to forge data format
@@ -241,6 +254,7 @@ def verify_module_for_inputs_deprecated(
         verify_module_for_inputs_deprecated(
             model=model,
             inputs=inputs,
+            compiler_cfg=compiler_cfg,
             pcc=pcc,
             dev_data_format=dev_data_format,
             convert_to_forge=convert_to_forge,
@@ -251,6 +265,7 @@ def verify_module_for_inputs(
         cls,
         model: Module,
         inputs: List[torch.Tensor],
+        compiler_cfg: CompilerConfig,
         verify_config: Optional[VerifyConfig] = VerifyConfig(),
         dev_data_format: forge.DataFormat = None,
         convert_to_forge: bool = True,  # explicit conversion to forge data format
@@ -259,6 +274,7 @@ def verify_module_for_inputs(
         verify_module_for_inputs(
             model=model,
             inputs=inputs,
+            compiler_cfg=compiler_cfg,
             verify_config=verify_config,
             dev_data_format=dev_data_format,
             convert_to_forge=convert_to_forge,