Using PT2 compile in tandem with the functional API

username: andreibratu
cooked: <p><strong>Background:</strong><br>
Hi everyone! I’m having an awesome time with PT2 and I would like some advice on how I can further integrate <code>torch.compile</code> in my code. I’m doing some research in differential privacy, and thus I make heavy use of the per_sample gradients and the functional API provided by torch throughout the training loop.</p>
<p><strong>Problem:</strong><br>
Here’s my simplified training loop:</p>
<pre><code class="lang-python">import torch
from torch.func import grad, vmap, functional_call

cfg = {} # Large config object for configuring train loop
device = 'cuda:0'
train_loader = '...'
model = '...'.to(device)
optim = '...'
params = dict(model.named_parameters())
buffers = dict(model.named_buffers())
criterion = lambda input, target: torch.nn.functional.cross_entropy(
        input, target, reduction='mean'
)

def compute_loss(model, params, buffers, batch, target, loss_fn):
    batch = batch.unsqueeze(0)
    targets = target.unsqueeze(0)
    predictions = functional_call(model, {**params, **buffers}, batch)
    loss = loss_fn(predictions, targets)
    return loss

# Mark static arguments/ constants with None or
# the dimension along which the operation ought
# to vectorize i.e. batch and and ground truth
compute_per_sample_grads = vmap(
    grad(compute_loss, argnums=1),
    in_dims=(None, None, None, 0, 0, None),
    randomness="same",
)

# ....

for batch, target in train_loader:
    batch = batch.to(device)
    target = target.to(device)
    y_pred = functional_call(model, params, (batch,))
    per_sample_grads = compute_per_sample_grads(
        model, params, batch, target, criterion
    )
    for param, grad_sample in zip(params.values(), per_sample_grads.values()):
        if param.grad_sample is None:
            param.grad_sample = grad_sample.contiguous()
    # ....
    optim.step()

# ...
</code></pre>
<p>The behaviour I expected was to be able to wrap the model with torch.compile and then use functional_call, but I receive <code>RuntimeError: Cannot access data pointer of Tensor that doesn't have storage</code>. I have also attempted to compile the <code>compute_per_sample_grads</code> () or <code>compute_loss</code> instead, resulting in the same error.</p>
<p>I have found issues regarding <a href="https://github.com/pytorch/pytorch/issues/100320" rel="noopener nofollow ugc">compile+vmap</a> or <a href="https://github.com/pytorch/pytorch/issues/94851" rel="noopener nofollow ugc">compile+grad</a>, but I could not find a situation of using both. I did not want to raise an issue since I feel that I’m misusing the API atm.</p>
<p>Many thanks for your time! I would be happy to provide other artificats as needed e.g. exact ResNet9 architecture.</p>
<details>
<summary>
Error Log</summary>
<p>(venv) bratua@encephalon:~/sparse-dp$ python train.py<br>
Files already downloaded and verified<br>
Files already downloaded and verified<br>
Number of epochs: 113<br>
Noise multiplier: 1.24<br>
[2023-05-05 15:40:49,259][opacus.data_loader][WARNING] - Ignoring drop_last as it is not compatible with DPDataLoader.<br>
Started training run: 1024_BS-1_LR-0.0_MTM-1.24_NOISE-5000_STEPS-1024_BS-2_CLIP-True_DP<br>
Number of parameters: 417,114<br>
Epoch 1:   1%|▋                                                                                                                                    | 233/45000 [00:00&lt;00:44, 1001.41it/s][2023-05-05 15:40:51,043] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing forward<br>
[2023-05-05 15:40:51,214] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo done tracing forward (RETURN_VALUE)<br>
[2023-05-05 15:40:51,223] torch._dynamo.output_graph: [INFO] Step 2: calling compiler function debug_wrapper<br>
[2023-05-05 15:40:53,190] torch._inductor.compile_fx: [INFO] Step 3: torchinductor compiling FORWARDS graph 0<br>
[2023-05-05 15:40:54,472] torch._inductor.compile_fx: [INFO] Step 3: torchinductor done compiling FORWARDS graph 0<br>
[2023-05-05 15:40:54,473] torch._dynamo.output_graph: [INFO] Step 2: done compiler function debug_wrapper<br>
Error executing job with overrides: []<br>
Traceback (most recent call last):<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/utils.py”, line 547, in preserve_rng_state<br>
yield<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 814, in wrap_fx_proxy_cls<br>
example_value = wrap_to_fake_tensor_and_record(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 957, in wrap_to_fake_tensor_and_record<br>
fake_e = wrap_fake_exception(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/utils.py”, line 808, in wrap_fake_exception<br>
return fn()<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 958, in <br>
lambda: tx.fake_mode.from_tensor(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py”, line 1324, in from_tensor<br>
return self.fake_tensor_converter(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py”, line 314, in <strong>call</strong><br>
return self.from_real_tensor(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py”, line 272, in from_real_tensor<br>
out = self.meta_converter(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_subclasses/meta_utils.py”, line 502, in <strong>call</strong><br>
r = self.meta_tensor(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_subclasses/meta_utils.py”, line 275, in meta_tensor<br>
base = self.meta_tensor(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_subclasses/meta_utils.py”, line 381, in meta_tensor<br>
s = t.untyped_storage()<br>
NotImplementedError: Cannot access storage of TensorWrapper</p>
<p>During handling of the above exception, another exception occurred:</p>
<p>Traceback (most recent call last):<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 324, in _compile<br>
out_code = transform_code_object(code, transform)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/bytecode_transformation.py”, line 445, in transform_code_object<br>
transformations(instructions, code_options)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 299, in transform<br>
tracer = InstructionTranslator(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 1670, in <strong>init</strong><br>
self.symbolic_locals = collections.OrderedDict(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py”, line 1673, in <br>
VariableBuilder(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 172, in <strong>call</strong><br>
return self._wrap(value).clone(**self.options())<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 238, in _wrap<br>
return self.wrap_tensor(value)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 639, in wrap_tensor<br>
tensor_variable = wrap_fx_proxy(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 754, in wrap_fx_proxy<br>
return wrap_fx_proxy_cls(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/variables/builder.py”, line 814, in wrap_fx_proxy_cls<br>
example_value = wrap_to_fake_tensor_and_record(<br>
File “/usr/lib/python3.8/contextlib.py”, line 131, in <strong>exit</strong><br>
self.gen.throw(type, value, traceback)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/utils.py”, line 549, in preserve_rng_state<br>
torch.random.set_rng_state(rng)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/random.py”, line 18, in set_rng_state<br>
default_generator.set_state(new_state)<br>
RuntimeError: Cannot access data pointer of Tensor that doesn’t have storage</p>
<p>Set torch._dynamo.config.verbose=True for more information</p>
<p>You can suppress this exception and fall back to eager by setting:<br>
torch._dynamo.config.suppress_errors = True</p>
<p>The above exception was the direct cause of the following exception:</p>
<p>Traceback (most recent call last):<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 104, in _fn<br>
return fn(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 262, in _convert_frame_assert<br>
return _compile(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/utils.py”, line 163, in time_wrapper<br>
r = func(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 394, in _compile<br>
raise InternalTorchDynamoError() from e<br>
torch._dynamo.exc.InternalTorchDynamoError</p>
<p>During handling of the above exception, another exception occurred:</p>
<p>Traceback (most recent call last):<br>
File “train.py”, line 185, in main<br>
per_sample_grads = compute_per_sample_grads(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 434, in wrapped<br>
return _flat_vmap(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 39, in fn<br>
return f(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 619, in _flat_vmap<br>
batched_outputs = func(*batched_inputs, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py”, line 1380, in wrapper<br>
results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 39, in fn<br>
return f(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py”, line 1245, in wrapper<br>
output = func(*args, **kwargs)<br>
File “train.py”, line 40, in compute_loss<br>
predictions = functional_call(model, {**params, **buffers}, batch)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_functorch/functional_call.py”, line 143, in functional_call<br>
return nn.utils.stateless._functional_call(<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/nn/utils/stateless.py”, line 262, in _functional_call<br>
return module(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1501, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 82, in forward<br>
return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 209, in _fn<br>
return fn(*args, **kwargs)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 337, in catch_errors<br>
return callback(frame, cache_size, hooks)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 404, in _convert_frame<br>
result = inner_convert(frame, cache_size, hooks)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py”, line 109, in _fn<br>
torch.cuda.set_rng_state(cuda_rng_state)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/cuda/random.py”, line 64, in set_rng_state<br>
_lazy_call(cb)<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/cuda/<strong>init</strong>.py”, line 183, in _lazy_call<br>
callable()<br>
File “/home/bratua/sparse-dp/venv/lib/python3.8/site-packages/torch/cuda/random.py”, line 62, in cb<br>
default_generator.set_state(new_state_copy)<br>
RuntimeError: Cannot access data pointer of Tensor that doesn’t have storage</p>
<p>Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.</p>
</details>