[RFC] Add HLO Minification API

[awskila]

🚀Feature

Add support for HLO minification in PyTorch/XLA similar to the PyTorch minifier.

Motivation

Currently (as of August 15th), there are three models that result in precision failures in the Dynamo benchmarks with the torch_xla_trace_once backend. All three are seen on BartForConditionalGeneration models. These errors are difficult to debug, as it’s difficult to narrow down the operator(s) that cause the failure.

BartForConditionalGeneration: [2023-08-04 00:10:21,049] torch._dynamo.utils: [ERROR] RMSE (res-fp64): 0.40076, (ref-fp64): 0.00000 and shape=torch.Size([1, 1024, 50265])

MBartForConditionalGeneration: [2023-08-04 00:57:55,593] torch._dynamo.utils: [ERROR] RMSE (res-fp64): 0.54682, (ref-fp64): 0.00000 and shape=torch.Size([1, 1024, 50265])

PLBartForConditionalGeneration: [2023-08-04 01:26:40,250] torch._dynamo.utils: [ERROR] RMSE (res-fp64): 0.71820, (ref-fp64): 0.00000 and shape=torch.Size([1, 1024, 50005])

These three models pass on Inductor backend and on Eager mode. It is easier to debug model evaluation failures on those two modes, as the PyTorch minifier can be used. It produces a minified subgraph to narrow down the operator(s) causing the failure. The PT minifier currently only works for FX IR minification. Thus, we need a minifier at the HLO level.

We are also interested in an HLO minifier because it would help AWS Neuron customers. From the customer’s perspective, it is usually not easy for the customers to share the entire model when encountering compilation/precision errors. If they’re able to run the minifier on their model, they can share the minified subgraph with us while protecting their IP. We will also be able to more quickly debug issues, leading to an improved customer experience.

Background on Existing PyTorch Minifier

Usages

The PyTorch minifier is a TorchDynamo debugging tool that narrows down the FX graph to its minimal subgraph that reproduces a given precision or compilation error. The forwards and backwards graphs are traced in TorchDyanmo and AOTAutograd, and these are used as inputs for the minifier. The output is a minimal subgraph that may include as few as one layer, thus greatly simplifying the debugging process.

In the case of a compilation failure, the minifier function is called right after the FX graph is traced. For a precision failure, during the execution of an eval loop, the minifier is ran upon an instruction having a root-mean-square error greater than the tolerance. The minifier is triggered by setting two environment variables - TORCHDYNAMO_REPRO_AFTER, which specifies whether to run the minifier on forwards or backwards graphs in case if compilation fails, and TORCHDYNAMO_REPRO_LEVEL, which specifies the granularity of the minifier - 1 specifies a compilation error, and 4 specifies a precision error.

Upon finding the minimal possible subgraph, it saves it as a Repro class method (usually defined to a reproducer script, repro.py. This script invokes TorchDynamo using torch._dynamo.optimize and passes the minified subgraph and the name of the compilation backend (e.g. aot_eager). Finally, a forward pass (and a backward pass, if AOTAutograd) is ran on the minimal subgraph, and the error is reproduced.

Current Implementation

The PyTorch minifier minimizes an FX graph with a given set of inputs. It attempts four strategies to minimize the graph, include suffix truncation, delta debugging, eliminating dead code, and removing unused inputs. It attempts these strategies until we’re left with the minimal possible subgraph that can reproduce the compilation or precision failure.

Suffix Truncation

The primary method used to minify the subgraph is suffix truncation. This is performed by trying to remove latter parts of the graph using a binary-search-like algorithm. It stops this when we can no longer reproduce the failure, and returns the subgraph. A step-by-step example on how it works can be seen in the Appendix section.

Delta Debugging

If the suffix truncation strategy is not sufficient, such as in cases where there are lots of prefix layers in the graph, delta debugging will be used as the secondary minification method. This strategy elevates prefix layers as input arguments to the forward function. An example of delta debugging is seen in the Appendix section.

Pitch

We want to propose an HLO minifier that mimics the behavior of existing FX minifier. It takes an HLO graph as input, and returns a minified HLO subgraph as the output. The minification tool can also generate a script that includes the subgraph. Such a script enables easier error reproduction and fix validation. This tool will be available for PT/XLA at first. However, the wider StableHLO community can also benefit from this minifier later. The proposed minification tool should be able to be easily integrated with other workflows such as TorchDynamo. More specifically, we will be using TorchDynamo as an entry point for testing the POC.

Implementation

Our intention is for this tool to be used in PyTorch, TensorFlow, JAX, and other frameworks. Thus, most of the new code we propose will be implemented at the C-level, including the HLO parsing/interpretation and the minification strategies. The code will be ported to Python via bindings, similar to current functions such as _get_xla_tensors_hlo. The final minified subgraph can be passed into a framework-specific codegen, creating a short Python script to reproduce the compilation or precision failure. This script can be then be ran to reproduce the error.

To implement the HLO minifier, the following code changes will need to be made to the PyTorch/XLA repo.

1. In the torch_xla/debug directory, create an hlo-minifier script. This will be imported by the user from a trainer script, and will support both compilation and precision minification use-cases.

2. Leverage CreateModuleFromProto, ParseHloModule, and ParseOperands functions to pull the list of modules from the HLO protobuf objects and parse the modules to get the list of ops/operands.

   • This will be done at the C level, with the code stored in the torch_xla/csrc directory.

3. Create functions that perform suffix truncation and delta debugging minification strategies, and also auxiliary functions that eliminate unused inputs/dead code.

   • Suffix truncation and delta debugging will be implemented at the C level, and will be imported to Python through bindings. These strategies will modify the HLO graph. The new functions will be defined in the init_python_bindings.cpp file.
   • The function that takes care of cleaning unused inputs and dead code will be implemented at the Python level, as it does not require any testing or analysis of the HLO graph.
   • Then, call these minification functions in a runner called run_hlo_minifier.

4. Build and run the minified graph to try and reproduce the error. This would be done at the Python level, and would differ based on the compiler used.

5. Repeat steps 3 and 4. Iteratively apply the minification strategies until the smallest possible subgraph that can reproduce the error is created.

6. Create a Python function that generates the repro.py script if a model fails compilation or precision.

We will add the HLO minifier code under torch_xla/debug. This Python module will be imported from PyTorch/XLA in order to invoke the HLO minifier from a trainer script. Finally, the run_hlo_minifier function would be called from a trainer script defined by the user.

API

A rough example is shown below to demonstrate the API and call flow. Let’s assume that torch.cat is causing a compilation error when attempting to concatenate two tensors with different shapes across a common dimension. The example uses torch.compile and thus is dependent on TorchDynamo, but our intention is for the HLO minifier to work with other compiler frontends. Please note that this code is high-level and its only purpose is to open up a discussion.

Trainer script

The trainer script is provided by the user. They import the HLO minification runner function from torch_xla.debug. The user should specify “compilation” or “precision” to force the HLO minifier to test for a specific error - otherwise, it will default to check for precision errors.

from torch.nn import *
from torch_xla.debug import run_hlo_minifier

class Model(torch.nn.Module):
   def __init__(self):
        super().__init__()
        self.self_mlp_0 = Linear(in_features=64, out_features=256, bias=True)

    def forward(self, cat):
        self_mlp_0 = self.self_mlp_0(cat)
        x = torch.add(torch.randn(4), torch.ones(4))
        y = torch.mul(torch.randn(8,4), torch.randn(8,4))
        cond = torch.randn(3,4) 
        cat_1 = torch.cat((y, cond), dim = 0) #assume that cat is causing an error
        return (cat_1, self_mlp_0)

inputs = [torch.randn(3), torch.randn(4)]
model = Model()

try:
    trained_model = torch.compile(model, backend="torchxla_trace_once")
    run_trained_model = trained_model(inputs)
except Exception as e:
    print(f"Encountered {type(e)}: {e}")
    run_hlo_minifier(compiled_model, inputs, "compilation")

baseline_model = torch.compile(model) #eager mode baseline
baseline_trained_model = baseline_model(inputs)

#run minifier if delta is observed vs baseline
if not torch.allclose(run_trained_model, baseline_trained_model):
    run_hlo_minifier(trained_model, inputs, "precision")
else:
    print(f"Model passes evaluation")

Output: Repro script

The repro script includes the minified subgraph and the user-defined inputs. If the error is precision-related, a comparison against a golden generated with eager mode will be added to the repro script. Our example is a compilation error, so we skip this.

from torch.nn import *

class MinifiedModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
     
    def forward(self, y, cond):
        return torch.cat((y, cond), dim = 0)

inputs = [torch.randn(3), torch.randn(4)]
model = MinifiedModel()
optim_model = torch.compile(model, backend="torchxla_trace_once")
optim_model(inputs)

HLO minifier script

The HLO minifier runner function, minifier strategy functions, and other helper function prototypes are displayed below. This code will be saved in the torch_xla/debug directory.

import torch_xla

def run_hlo_minifier(model, inputs, error_type="precision"):
   #Iterate through all HLO modules, cycling through minification strategies and create reproducer
   modules = _suffix_truncation(model.modules(), inputs, error_type)
   modules, input_args = _delta_debugging(modules, inputs, error_type)
   input_args = clean_up_unused_args(modules, input_args)
   create_repro_script(modules, input_args, inputs)

def _suffix_truncation(modules, inputs, error_type="precision"):
    # Try and truncate suffix using pre-set seeds, mimicking the behavior of FX minifier
    # Attempt to truncate HLO modules
    modules = torch_xla._XLAC._suffix_truncation(modules, inputs, iter, error_type)
    
    # Verify we can reproduce the error. Call suffix truncation function as needed
    # Increment iterator until we cycle through all seeds, or when graph has 1 node
    ...
    return modules

def _delta_debugging(modules, inputs, error_type="precision"):
    # Try and remove prefix, and elevate variables to input arguments for the graph
    mod_modules, inputs = torch_xla._XLAC._delta_debugging(modules, inputs, iter, error_type)
 
    # Verify we can reproduce the error. Call delta debugging function as needed
    # Increment iterator until we reach the end of HLO modules
    ...
    return (mod_modules, inputs)
         
def clean_up_unusued args(modules, input_args):
    #Try and clean up dead code by looking for unused input args
    ...
    return mod_input_args

def create_repro_script(minified_modules, input_args, inputs):
    with ("repro.py", 'w') as repro_script:
        # write repro script that calls minifed graph to reproduce the model error
        ...

Additional context

After the above is implemented, we are interested in modifying minifier functionality in an attempt to improve its granularity. This can be done by using graph pruning and model compression techniques. This would be applied to both FX and HLO minifiers, with the goal of trying to narrow down the minified graph to a single instruction more reliably. We want to build on existing open minifier issues regarding accuracy minification and help with some issues seen with handling quantization and dynamic shape support use-cases.

Appendix

Suffix Truncation Example (FX minifier)

Let’s say that our backend fails on division ops. Let’s say the forward function looks something like this:

def forward(a_in, b_in):
   b_in = torch.ops.aten.div.Tensor(a_in, b_in)
   a_in = torch.ops.aten.add.Tensor(a_in, b_in)
   b_in = torch.ops.aten.mul.Tensor(a_in, b_in)
   a_in = torch.ops.aten.sqrt(b_in)
   a_in = torch.ops.aten.add.Tensor(b_in, b_in)
   b_in = torch.ops.aten.sigmoid(a_in)
   b_in = torch.ops.aten.maximum(a_in, b_in)
   a_in = torch.ops.aten.pow.Tensor_Tensor(b_in, a_in)

We first cut the last half of the graph.

def forward(a_in, b_in):
   b_in = torch.ops.aten.div.Tensor(a_in, b_in)
   a_in = torch.ops.aten.add.Tensor(a_in, b_in)
   b_in = torch.ops.aten.mul.Tensor(a_in, b_in)
   a_in = torch.ops.aten.sqrt(b_in)

We can still reproduce the error, so we cut the last 3/4ths of the graph.

def forward(a_in, b_in):
   b_in = torch.ops.aten.div.Tensor(a_in, b_in)

We can still reproduce the error, and we’re left with a minified subgraph.

Delta Debugging + Clean-Up Example (FX minifier)

Let’s say that our backend fails on multiplication ops. For example, let’s say our forward function looks like this:

def forward(a_in):
  b_in = x / 10
  c_in = b_in - 5
  d_in = c_in * 5
  return d_in

If a multiplication operation causes the graph to fail, the middle node can be removed, with c being promoted to an input.

def forward(a_in, c_in):
  b_in = x / 10
  d_in = c_in * 5
  return d_in

Similarly, we can do this for b_in and x .

def forward(a_in, c_in, b_in, x):
  d_in = c_in * 5
  return d_in

This is close, but the graph isn’t fully minified at this point. This is because we have unused inputs. There are two minor, trivial remaining strategies used, which are to clear unused inputs and any dead code. We don’t need a_in, b_in, or x to reproduce the failure in our 1-layer minified graph, so let’s clear them. Also, this function can become a one-line function by moving d_in to the return statement.

def forward(c_in):
  return c_in * 5

Now the graph is fully minified using delta debugging, and cleaning up unused inputs and dead code.

cc @alanwaketan, @JackCaoG, @qihqi

[GleasonK]

After a little bit of digging I found the following reproducer / minifiers:

• hlo_bisect looks promising for this use case but uses different minification strategy
• mlir_bisect for mlir-interpreter
• mlir-reduce, could likely build something on top of StableHLO using this.

I like the notion of having one tool for all frameworks to target, depending on the use cases perhaps such a tool should be written in StableHLO against the reference interpreter (i.e. if for conversion bugs, StableHLO makes sense, for compiler bugs HLO is better).

Do you have any ideas for how the python reproducer part would work for an HLO-based tool? Is the expectation for there to be a way to go from minified HLO --> PyTorch model?

The PT minifier currently only works for FX IR minification. Thus, we need a minifier at the HLO level.

Perhaps a naive question, but is there no way to go from FX IR --> HLO? APIs like save_as_stablehlo make me think this is possible, is this just for inference maybe? Could the existing tool be used with some new PT/XLA APIs? (cc @qihqi) I'm guessing theres a bit of nuance I'm missing here, but would be good for someone from PT/XLA side to evaluate feasibility of using the existing minifier as well.

[awskila]

Thanks @GleasonK for finding the existing minifier-like tools and for the insight! The bisect and reduction tools looks interesting - I wanted to the tool to be compatible for StableHLO and MLIR-HLO, so maybe it's a matter of evaluating and potentially adding additional minification strategies onto the existing HLO bisect tool. Some comments inline:

I like the notion of having one tool for all frameworks to target, depending on the use cases perhaps such a tool should be written in StableHLO against the reference interpreter (i.e. if for conversion bugs, StableHLO makes sense, for compiler bugs HLO is better).

Yes, I like this idea too. One problem with the existing FX minifier is that it's written only for a torch.compile and lacks any HLO support. This tends to cause problems for anyone working on custom silicon. A framework-agnostic minifier that takes in the HLO graph as input is definitely goal here. From the framework itself, it would just be a matter of adding the library header from PT/XLA and calling a single minifier function that would call the various minification functions and return the minified graph.

Do you have any ideas for how the python reproducer part would work for an HLO-based tool? Is the expectation for there to be a way to go from minified HLO --> PyTorch model?

That's a very good question. We'd need to use a parser to pull the parameters from HLO IR, and then a framework-level interpreter to convert that to the intended Python framework code. I was thinking to support PyTorch at first, but to also add support for a similar TensorFlow and JAX. The reproducer script for the FX minifier is really convenient, and we should try and include it somehow. It'll be a bit of a challenge for an HLO-based tool. For the FX minifier, given that tool is entirely Python-based and leverages TorchDynamo, it's easier to create such a script.

Perhaps a naive question, but is there no way to go from FX IR --> HLO? APIs like save_as_stablehlo make me think this is possible, is this just for inference maybe? Could the existing tool be used with some new PT/XLA APIs?

Oh wow. This looks very interesting - it appears to leverage torch.export as well. I see one potential bottleneck - it's possible that certain features like dynamic shape support and quantization may not be translated from FX IR to HLO IR successfully, and maybe certain operator lowerings could be invalid. These aren't minifier-specific problems though. Other than that, I don't see why the current tool could not be extended, though I'm merely a user of the minifier.

I'm guessing theres a bit of nuance I'm missing here, but would be good for someone from PT/XLA side to evaluate feasibility of using the existing minifier as well.

Agreed - I would love to hear from someone from the PT/XLA team to help determine feasibility. If possible, I would love to hear from someone more familiar with the current state of torch.export, as it's fairly new and has caused some problems for me and other users for certain Hugging Face models. If exporting an FX IR model that was built on an XLA backend (i.e. torchxla_trace_once) works reliably, it seems like this could be the best option to pursue. We'd still have to handle producing the repro script, but potentially the minification can be done on the FX side, and we could port the minified FX model over to HLO? That would be quite nice - obviously it relies on PyTorch to be used to produce the exported model, but still quite nice.

[awskila]

So I took a look at the feasibility of porting the FX IR minifier to HLO. The minifier takes an fx.GraphModule object as an argument, and returns an fx.GraphModule object.

Then torch.export.export needs a torch.nn.Module object, and this function returns an ExportedProgram object. The save_as_stable_hlo function needs an ExportedProgram object.

While you can use fx.symbolic_trace to convert a torch.nn.module object to an fx.GraphModule object, converting an fx.GraphModule object back to a torch.nn.module object leads to problems (as far as I'm aware). So it looks like using the FX minifier and exporting the minified FX graph to Stable HLO doesn't look feasible, unless I'm missing something here.

[awskila]

Thanks a bunch @GleasonK for pointing me towards the recently added Torch-to-StableHLO porting functions! Have opened a new PR that ports the minified FX graph to HLO. Take a look at it when you get a chance.

If this is sufficient, we can close this issue for now after the PR is merged. Porting the minified FX graph to HLO is probably the cleanest/simplest solution to this problem.

[qihqi]

Nit:

converting an fx.GraphModule object back to a torch.nn.module object leads to problems (as far as I'm aware).

fx.GraphModule is a nn.Module subclass and can be used where nn.Module can, without any conversion.

If the tool you are looking is outputting HLO subgraph, then as Kevin mentioned, there are 2 ways: one is dump full HLO graph from torch; then run mlir-bisect; another is to run fx minifier with dynamo with torchxla backend; and once we have minified fx graph, we use torchxla to convert it to stablehlo. Both should be equally feasible.

[awskila]

fx.GraphModule is a nn.Module subclass and can be used where nn.Module can, without any conversion.

Thanks @qihqi for the response. I came to that realization late last week - had thought that there would be some layers of the network would be lost if casting an fx.GraphModule object to a torch.nn.module object. Was definitely wrong about that.

If the tool you are looking is outputting HLO subgraph, then as Kevin mentioned, there are 2 ways: one is dump full HLO graph from torch; then run mlir-bisect; another is to run fx minifier with dynamo with torchxla backend; and once we have minified fx graph, we use torchxla to convert it to stablehlo. Both should be equally feasible.

Yep, that's definitely the goal. I have a PR (I just CC'd you to it) that takes the minified FX graph, and uses the conversion function that you and your team added. It's a debug tool and obviously is dependent on torch_xla, so I have this functionality gated via the XLA_HLO_DEBUG flag.

Also, I really must thank you and the rest of the StableHLO team for implementing this feature. My hlo-minifier-from-scratch library was having some problems when I was testing it on various models. Was stuck in a seemingly never-ending debugging process over the past ~2 weeks. Using the existing minifier and porting it to stablehlo is a lot cleaner and elegant.

[GleasonK]

This is great! Glad it ended up working out.

Once that PR lands would be great to get a simple usage test in the PT/XLA repo, and maybe a brief documentation in TROUBLESHOOTING.md. Will leave that to PT/XLA folks to determine how to go about that.

[awskila]

Me too, thanks again for the help! And yes, good suggestion to add an end-to-end test and document this feature.

[awskila]

Closing the issue as the PR has been merged to the PyTorch repo. Thanks everyone for the feedback and help!