How should I use torch.compile properly?

username: philgzl
cooked: <p>I have multiple questions about how to use torch.compile properly. I found the tutorial/documentation lackluster. I will try and list all of them down including those I found answered in this forum but are missing from the tutorial, for future readers.</p>
<ul>
<li>
<p>How to use torch.compile with a non-trivial nn.Module? <a href="https://stackoverflow.com/questions/75886125/how-should-i-using-torch-compile-properly" rel="noopener nofollow ugc">This post on stackoverflow</a> perfectly sums up my question. Namely, should torch.compile be called manually on every sub-module I wish to compile, or does it handle that automatically? More generally, how does torch.compile behave when used on a nn.Module? Does it compile only the forward method? In my case, similar to OP in that stackoverflow post, I observe zero performance difference when using torch.compile, even though my model has 27M parameters and I am using one of the recommended GPUs in the tutorial (V100). I also note torch.compile exits immediately, as if it were doing nothing.</p>
</li>
<li>
<p>Should I cast to DDP before or after torch.compile? In the DistributedDataParallel tutorial (can’t post another link due to new forum user limitation), torch.compile is called after on the DDP model. But <a href="https://discuss.pytorch.org/t/torch-compile-before-or-after-cuda/176031">in this post</a>, it is recommended to call torch.compile before. So what is the right answer?</p>
</li>
<li>
<p>Should I move the model to device before or after torch.compile? In that same post, it is recommended to cast to CUDA before calling torch.compile. But this is not mentioned in the tutorial even though this seems like an pretty essential consideration.</p>
</li>
</ul>

username: marksaroufim
cooked: <p>Hi</p>
<ul>
<li>I answered the question on stack overflow, one aspect that’s important to note is that V100 improvements are alright for torch.compile but the real speedups will come from A100 or A10G. And yes indeed nothing will happen when you torch.compile only, the compilation will happen at the time of the first inference. The name is sorta bad, it should be called torch.jit but that was already taken XD</li>
<li>Regarding your question on distributed, today the distributed support is not very fleshed. There’s a tradeoff if you compile the DDP module then torch.compile should be able to trace the communication and do more optimizations there but it doesnt today so stuff is likely to break so it’s safer to compile the inner module but this will likely evolve in our next releases</li>
<li>It’s not a huge deal either way, inductor would just prefer it if you don’t have too many device copies but that’s not gonna break it or cripple perf</li>
</ul>

username: philgzl
cooked: <p>Thank you so much for the quick reply!</p>
<p>I have some follow-up questions. I have a trainer class that takes an initialized nn.Module as argument. The model defines <code>train_step</code> and <code>val_step</code> methods that are called by the trainer during training, à la Lightning. I want to be able to pass a <code>compile</code> boolean argument to my trainer to call <code>torch.compile</code> on the model. Something like this (code heavily simplified for this question):</p>
<pre><code class="lang-python">class MyTrainer:
    def __init__(self, model, compile, ...):
        if compile:
            model = torch.compile(model)
        self.model = model
        ...  # initialize dataloaders

    def run(self):
        for epoch in range(self.epochs):
            self.model.train()
            for batch in self.train_dataloader:
                train_loss = self.model.train_step(batch)
            ...  # accumulate and log train loss
            self.model.eval()
            for batch in self.val_dataloader:
                val_loss = self.model.val_step(batch)
            ...  # accumulate and log val loss


class MyModel:
    def __init__(self, ...):
        ...  # initialize model, optimizer and criterion

    def train_step(self, batch):
        self.optimizer.zero_grad()
        loss = self._step(batch)
        loss.backward()
        self.optimizer.step()
        return loss

    def val_step(self, batch):
        loss = self._step(batch)
        return loss

    def _step(self, batch):
        inputs, labels = batch[:, 0], batch[:, 1]
        outputs = self(inputs)
        loss = self.criterion(outputs, labels)
        return loss.mean()
</code></pre>
<p>This currently does nothing. I believe it’s because by calling <code>model = torch.compile(model)</code> and assigning <code>self.model = model</code>, I lose reference to the original model and the function to compile thus becomes <code>self.model.__call__</code> instead of the original model’s <code>__call__</code>. So when calling <code>train_step</code> and <code>val_step</code>, as these fall back to the original model’s implementation, they use the un-optimized model’s <code>__call__</code>. Am I correct?</p>
<p>I found out that using <code>torch.compile</code> this way instead does do something. Namely I compile the <code>forward</code> method manually:</p>
<pre><code class="lang-python">class MyTrainer:
    def __init__(self, model, compile, ...):
        if compile:
            model.forward = torch.compile(model.forward)
        self.model = model
        ...  # initialize dataloaders
</code></pre>
<p>Alternatively compiling the <code>train_step</code> and <code>val_step</code> methods this way also does something, though the result is different as per the <code>torch_dynamo</code> logs:</p>
<pre><code class="lang-python">class MyTrainer:
    def __init__(self, model, compile, ...):
        if compile:
            model.train_step = torch.compile(model.train_step)
            model.val_step = torch.compile(model.val_step)
        self.model = model
        ...  # initialize dataloaders
</code></pre>
<p>Are these valid approaches for my use case? Should one be preferred over the other? Note that in the second solution, the compiled function contains the <code>optimizer.zero_grad()</code>, <code>loss.backward()</code> and <code>optimizer.step()</code> calls. Is it wrong to compile these?</p>
<p>On another note, from what I understand from the logs produced by <code>torch._dynamo</code>, compilation happens not only on the first batch, but on every batch during the first epoch. Is this normal? Starting the second epochs, no more logs are produced.</p>

username: marksaroufim
cooked: <p>Yeah integrations with training loop providers can be funky, this is a different take on the API <a href="https://github.com/pytorch/pytorch/pull/97565" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/pull/97565</a> - lmk what you think because i tried to include in core exactly what you’ve done</p>

username: philgzl
cooked: <p>Mmh I see, so if I understand correctly from <a href="https://github.com/msaroufim/mlsys-experiments/blob/main/save-compiled-model.py" rel="noopener nofollow ugc">here</a>,</p>
<ul>
<li>You define an in-place <code>compile</code> method for the model which assigns the optimized module to an attribute of the model</li>
<li>When calling the model, you check if the optimized module attribute has been assigned and use it if it was</li>
<li>When saving the model, you make sure the optimized model state is not dumped</li>
<li>If <code>persistent_compilation</code> is <code>True</code>, the model is compiled upon loading</li>
</ul>
<p>Is this correct? I can foresee this can be easily applied to my framework by extended my model base class. Will report back.</p>

username: philgzl
cooked: <p>So I tried implementing the idea in your script linked above as follows:</p>
<pre><code class="lang-auto">class MyBaseModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.optimized_module = None

    def compile(self, *args, **kwargs):
        self._optimized_module = torch.compile(self, *args, **kwargs)

    def __call__(self, *args, **kwargs):
        if self._optimized_module:
            return self._optimized_module(*args, **kwargs)
        else:
            return super().__call__(*args, **kwargs)

    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove _optimized_module from the state
        state.pop('_optimized_module', None)
        return state
</code></pre>
<p>However this produces a <code>RecursionError: maximum recursion depth exceeded</code> error when calling <code>model.train()</code> and probably other methods that iterate over child modules. This I guess makes sense since eventually it calls <code>self._optimized_module.train()</code> which falls back to the un-optimized module’s <code>.train()</code> and thus produces the infinite recursion.</p>
<p>The following, which is a subset of the PR you linked, prevents the recursion and seems to work:</p>
<pre><code class="lang-auto">class MyBaseModel(nn.Module):
    def __init__(self):
        super().__init__()
        self._compiled_call_impl = None

    def compile(self, *args, **kwargs):
        self._compiled_call_impl = torch.compile(self._call_impl, *args, **kwargs)

    def __call__(self, *args, **kwargs):
        if self._compiled_call_impl is not None:
            return self._compiled_call_impl(*args, **kwargs)
        else:
            return self._call_impl(*args, **kwargs)

    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove _compiled_call_impl from the state
        state.pop('_compiled_call_impl', None)
        return state
</code></pre>
<p>Should I be wary of using internal variables like <code>_call_impl</code>? Also, is the rest of the code in the PR necessary for what I want to achieve?</p>