How to solve the graph break happen in torch.compile

username: Ghogha_Atif
cooked: <p>I am solving the problem of my fine tuning a model am i am using a qlora to fine tune my model and using a torch.compile but i am getting a graph break in my line.how can i solve this what is the possible way anyone guide me for this i have learned the torch.compile is not support the 4 bit quantization how to disable or remove it from the compilation in mlp or in attention layer.</p>
<p>torch_compile_options = {<br>
“epilogue_fusion”   : True,<br>
“max_autotune”      : True,<br>
“shape_padding”     : True,<br>
“trace.enabled”     : True,<br>
“triton.cudagraphs” : False,<br>
}</p>
<h1><a name="p-465828-torchcompilefullgraphfalsedynamictrueoptionstorch_compile_options-1" class="anchor" href="#p-465828-torchcompilefullgraphfalsedynamictrueoptionstorch_compile_options-1"></a><span class="mention">@torch.compile</span>(fullgraph=False,dynamic=True,options=torch_compile_options)</h1>
<h1><a name="p-465828-def-compiled_llama_mlpselfx-2" class="anchor" href="#p-465828-def-compiled_llama_mlpselfx-2"></a>def compiled_llama_mlp(self,x):</h1>
<h1><a name="p-465828-down_projselfdown_projselfact_fnselfgate_projx-selfup_projx-3" class="anchor" href="#p-465828-down_projselfdown_projselfact_fnselfgate_projx-selfup_projx-3"></a>down_proj=self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))</h1>
<h1><a name="p-465828-return-down_proj-4" class="anchor" href="#p-465828-return-down_proj-4"></a>return down_proj</h1>
<h1><a name="p-465828-transformersmodelsllamamodeling_llamallamamlpforward-compiled_llama_mlp-5" class="anchor" href="#p-465828-transformersmodelsllamamodeling_llamallamamlpforward-compiled_llama_mlp-5"></a>transformers.models.llama.modeling_llama.LlamaMLP.forward = compiled_llama_mlp</h1>
<p>and from training i am using the hugging face tranier not creating a traning script from scratch</p>

username: RAMESH_BABU
cooked: <p>I’m currently encountering the same issue. Torch Dynamo doesn’t support several operations used in the BitsAndBytes (BnB) quantization process. One approach is to modify BnB so that almost every operation becomes traceable by Dynamo; alternatively, we could update Dynamo to support the operations BnB relies on.</p>
<p>For instance, some of the unsupported operations include:</p>
<ol>
<li>In the Linear4bit module, the Params4bit object contains a transpose operation on a tensor within a user-defined object, which isn’t supported.</li>
<li>Calling <strong>A.data_ptr()</strong> on a tensor isn’t supported.</li>
<li>Using <strong>ctypes.c_void_p</strong> variables is also not supported by Dynamo.</li>
</ol>
<p>While a simple workaround would be to disable Dynamo for these operations and revert to eager mode, that defeats our goal of eliminating graph breaks. Since our aim is to stitch everything into a single graph, the conclusion I’ve reached is that we need to modify some of the parts in <strong>BnB</strong> and <strong>torch._dynamo</strong> to support these quantization operations at all graph breaks.</p>
<p>Fun fact: The native model doesn’t trigger any graph breaks, LoRA causes just a single, easily fixable graph break, whereas the QLoRA model results in 39–40 graph breaks.</p>
<p>If anyone has a simpler solution, please share—I’ve already spent over two days fixing these graph break issues and still have 3–4 more to address (around 6–7 unique graph break reasons, which cumulatively result in 39–40 breaks).</p>

username: bdhirsh
cooked: <p>If the BnB code you’re calling is effectively a bunch of custom kernels (which sounds likely given the data pointer access), then the right approach here is generally to refactor bits-and-bytes’ kernels in to custom PyTorch operators, so they can be properly plumbed through the compiler. There’s a nice guide for it here: <a href="https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html" class="inline-onebox" rel="noopener nofollow ugc">PyTorch Custom Operators — PyTorch Tutorials 2.6.0+cu124 documentation</a></p>

username: niconunezz
cooked: <p>how did you fix the first one, i have a similar problem but nothing seems to work?</p>

username: RAMESH_BABU
cooked: <p>The issue affects every subclass of Tensor. Fortunately, the torch community has already resolved it in the latest nightly release. You can either switch to this nightly version for an immediate fix or examine its changes and patch them to your current stable version.</p>

username: niconunezz
cooked: <p>I’ve been trying to find the specific commit, but it appears nowhere. Do you have any clue?</p>

username: byi8220
cooked: <p>I think we ran into the same issue <img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"></p>
<p>Digging around in <a href="https://pytorch.org/docs/stable/torch.compiler_dynamo_deepdive.html#implementing-cpython-in-python" class="inline-onebox" rel="noopener nofollow ugc">Dynamo Deep-Dive — PyTorch 2.6 documentation</a> I noticed this,</p>
<blockquote>
<p>In particular, sometimes we find that an object is tracked as a <code>UserDefinedObjectVariable</code> (this is Dynamo’s catch-all class), when it should have been tracked as something more specific. In these cases, the <code>SourceBuilder.__call__</code> logic is often to blame.</p>
</blockquote>
<p>Would you agree with the idea that this is a torch compiler bug, where torch fails to properly wrap user defined subclasses of <code>Tensor</code> as <code>TensorVariable</code> (a bug which was fixed in nightly?)</p>
<p>Also,</p>
<blockquote>
<p>The issue affects every subclass of Tensor.</p>
</blockquote>
<p><em>Every</em> subclass, or just user defined subclasses? Is this actually breaking on torch’s own tensor subclasses, e.g. <code>torch.nn.Parameter</code>?</p>

username: RAMESH_BABU
cooked: <p>I believe the issue originates from the _wrap method in the VariableBuilder class. This bug does not affect PyTorch’s built-in tensor subclasses but specifically impacts user-defined classes that inherit from Tensor.</p>

username: byi8220
cooked: <blockquote>
<p>I believe the issue originates from the _wrap method in the VariableBuilder class.</p>
</blockquote>
<p>Hm, where do you think the issue may be and do you know if/how the nightly fixed it? I’m hesitant to just bump up to the latest nightly, since this may cause a ton of version breakage.</p>
<p>IIUC, we can explicitly mark a tensor subclass as traceable via <a href="https://github.com/pytorch/pytorch/blob/a3c6e3139a0f3418ae2c89df4b2b87d6dc97b0e2/torch/_dynamo/config.py#L156-L173" rel="noopener nofollow ugc">_dynamo.config.traceable_tensor_subclasses</a>, with heavy restrictions.</p>
<p>I suspect maybe one could register <code>Params4bit</code> here, and mess around with <code>__torch_forward__</code> and <code>__torch_dispatch__</code> for it to make something work. But I haven’t had any success there.</p>

username: RAMESH_BABU
cooked: <p>I was working on this issue for a small assignment without major dependencies. Since I was also busy with my semester exams at the time, I opted for a quicker fix by using the nightly version. I didn’t really look into how they actually solved it. <img src="https://discuss.pytorch.org/images/emoji/apple/cry.png?v=12" title=":cry:" class="emoji" alt=":cry:" loading="lazy" width="20" height="20"></p>

username: byi8220
cooked: <blockquote>
<p>I was working on this issue for a small assignment without major dependencies.</p>
</blockquote>
<p>I think we’re working on the same problem, and ran into the exact same issue. Upgrading to nightly caused my entire env to break, and I would have preferred a cleaner solution over hacking together a bunch of conflicting package versions. But for the sake of time, if it works it works.</p>

username: RAMESH_BABU
cooked: <p>If you’re working on Unsloth puzzles, you might find this repository helpful:</p>
<aside class="onebox githubrepo" data-onebox-src="https://github.com/RameshBabuAsh/UnslothPuzzles">
  <header class="source">

      <a href="https://github.com/RameshBabuAsh/UnslothPuzzles" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <div class="github-row" data-github-private-repo="false">
  <img width="690" height="344" src="https://opengraph.githubassets.com/e6566e6fc4064a1da3bcd5d2fa603bd8/RameshBabuAsh/UnslothPuzzles" class="thumbnail">

  <h3><a href="https://github.com/RameshBabuAsh/UnslothPuzzles" target="_blank" rel="noopener nofollow ugc">GitHub - RameshBabuAsh/UnslothPuzzles: This repository contains my solutions for Daniel...</a></h3>

    <p><span class="github-repo-description">This repository contains my solutions for Daniel Han's Unsloth challenge(five tasks) focused on optimizing deep learning workflows :)</span></p>
</div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>