Decomposition slows down lazy tensor tracing when using Adam optimizer

[lukeliu15]

🐛 Bug

Tracing time slows down significantly when using Adam optimizer.

This issue was originally reported here: https://dev-discuss.pytorch.org/t/decomposition-slows-down-the-lazy-tensor-tracing/1788

It seems that this torch-xla team started looking into this issue as of the 2.1 dev release notes: https://dev-discuss.pytorch.org/t/pytorch-xla-2-1-release-dev-update/1568#optimizer-tracing-time-regression-47
Are there any updates?

To Reproduce

Reproducible command can be found in the thread: https://dev-discuss.pytorch.org/t/decomposition-slows-down-the-lazy-tensor-tracing/1788/3

Expected behavior

No decomposition of aten::lerp_ and no tracing time regression.

Environment

• Reproducible on XLA backend [CPU/TPU/CUDA]: CUDA
• torch_xla version: 2.5.1

Additional context

Currently the workaround is to specify os.environ["XLA_DISABLE_FUNCTIONALIZATION"] = "1", but this flag is expected to go away and this flag currently breaks other parts of our model.

[ysiraichi]

Thank you for filing this issue.
I'd like to ask 2 things:

• Could you provide a small repro?
• Could you try running with the nightly version?

Without dynamo, I believe we shouldn't be using PyTorch decompositions. It should be going straight to the functionalization kernel. Maybe, what you are seeing is their meta functions running (fixed in pytorch/pytorch#136909)?

[lukeliu15]

Hey @ysiraichi, thanks for the response and the potential fix!

I verified the nightly version did not fix the issue.

Repro:

import os

os.environ["XLA_REGISTER_INSTALLED_PLUGINS"] = "1"
os.environ["XLA_DISABLE_FUNCTIONALIZATION"] = "1"

import torch
import torch_xla.core.xla_model as xm
import time

def simple_torch_model():
    layers = []
    layers.append(torch.nn.Linear(10, 100))
    layers.append(torch.nn.ReLU())

    for _ in range(98):
        layers.append(torch.nn.Linear(100, 100))
        layers.append(torch.nn.ReLU())

    layers.append(torch.nn.Linear(100, 1))

    return torch.nn.Sequential(*layers)

def train():
    device = xm.xla_device()
    model = simple_torch_model()
    model = model.to(device)

    optimizer = torch.optim.Adam(model.parameters())
    loss_fn = torch.nn.MSELoss()

    step_times = []

    for i in range(20):
        input = torch.randn(10).to(device)
        y = torch.randn(1).to(device)
        output = model(input)
        total_loss = loss_fn(output, y)
        optimizer.zero_grad()
        total_loss.backward()
        start_time = time.time()
        optimizer.step()
        optimizer_step_time = time.time() - start_time
        xm.mark_step()
        step_times.append(optimizer_step_time)
    
    median_time = sorted(step_times)[len(step_times)//2]
    print(f"Median optimizer step tracing time: {median_time:.4f}s")

Performance: (measured on A4000 machine)

	Median optimizer step tracing time
Base	0.48s
Include pytorch/pytorch#136909 fix	0.15s
Include os.environ["XLA_DISABLE_FUNCTIONALIZATION"] = "1" flag	0.04s

[ysiraichi]

Thank you for the profiling + repro.

I verified the nightly version did not fix the issue.

It seems that the second row improves performance, going from 0.08x slower FuncDisabled / Base to 0.26x FuncDisabled / PyTorchFix. Did you check whether the slowdown is still due to the decompositions being called? I expect that, for the second row, only the meta functions are being called (faster than the actual decompositions).

[lukeliu15]

Hey @ysiraichi, here are the stack trace results:

with pytorch/pytorch#136909 (this has not been merged to develop, right?):

disable functionalization:

the meta functions called are still quite slow, although faster than the decompositions.

[ysiraichi]

Well, I guess we could have a switch for disabling it inside the Functionalize generated kernels. The problem with this approach is that checks in the meta kernels wouldn't run. This could lead with harder to understand errors down the road. That said, I think this wouldn't be that big of a problem if you know there are no errors, and it leads to a big perf improvement.

@bdhirsh What do you think? Would you think an environment variable that disables meta kernels inside functionalization kernels be a good idea?

[bdhirsh]

It seems reasonable for XLA to have a config they can flip that tells the functionalization kernels not to run the meta reference implementations. FWIW, we effectively only run those meta reference kernels with XLA/LazyTensor today (since in torch.compile, we run a round of fake tensor prop once in dynamo before functionalization runs, which lets us catch any meta error then)

[ysiraichi]

What do you think of:

bool disableMetaReference() {
    static auto _value = std::getenv("TORCH_DISABLE_FUNCTIONALIZATION_META_REF");
    return _value != nullptr && strcmp(_value, "1") == 0;
}

at::Tensor & abs_(c10::DispatchKeySet dispatchKeySet, at::Tensor & self) {
  if (disableMetaReference()) {
    auto self_meta = to_meta(self);
    at::AutoDispatchSkipFunctionalize func_guard;
    c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch);
    at::_ops::abs_::call(self_meta);
  }
  ...
}

[bdhirsh]

hmm yeah, that seems ok (i see other places in aten where we directly read env vars from C++)