[torch-xla 2.1 - 2.4] when functionalization is on, there are no aliasing for gradients when using gradient accumulation

[jeffhataws]

🐛 Bug

When functionalization is on (XLA_DISABLE_FUNCTIONALIZATION=0), I see that there are fewer aliased tensors. Jack has a patch to increase the number of aliased tensors e3fc033 . However, even though this change helped increase the number of aliased tensor, it seems to still missing aliasing for gradients when gradient accumulation is used.

Using test_train_mp_mnist.py, make the modifications below. I added a mark_step to isolate the gradient accumulation loops.

@ -158,16 +163,19 @@ def train_mnist(flags, **kwargs):
       output = model(data)
       loss = loss_fn(output, target)
       loss.backward()
-      if flags.ddp:
-        optimizer.step()
-      else:
-        xm.optimizer_step(optimizer)
-      tracker.add(flags.batch_size)
-      if step % flags.log_steps == 0:
-        xm.add_step_closure(
-            _train_update,
-            args=(device, step, loss, tracker, epoch, writer),
-            run_async=flags.async_closures)
+
+      if step % 4 == 0:
+          xm.mark_step()
+          if flags.ddp:
+            optimizer.step()
+          else:
+            xm.optimizer_step(optimizer)
+          tracker.add(flags.batch_size)
+          if step % flags.log_steps == 0:
+            xm.add_step_closure(
+                _train_update,
+                args=(device, step, loss, tracker, epoch, writer),
+                run_async=flags.async_closures)

I only see 2 alias even though we expect all the gradient tensors to be aliased:

2024-06-03 21:15:37.676472: I torch_xla/csrc/xla_graph_executor.cpp:1462] Parameter sequence graph hash b8e15ed0391b82171706a34d84ca8ea0
2024-06-03 21:15:37.678822: I torch_xla/csrc/xla_graph_executor.cpp:1299] Aliased paramter 13 with output 4: s64[]
2024-06-03 21:15:37.678862: I torch_xla/csrc/xla_graph_executor.cpp:1299] Aliased paramter 14 with output 5: s64[]
2024-06-03 21:15:37.679222: I torch_xla/csrc/xla_graph_executor.cpp:1397] Compiling IR graph hash b8e15ed0391b82171706a34d84ca8ea0 on device CPU:0 ...

To Reproduce

Steps to reproduce the behavior:

1. Check out r2.1_aws_neuron branch
2. Apply a patch from Jack e3fc033
3. Build/install as in CONTRIBUTION doc
4. Go into xla/test
5. Edit test_train_mp_mnist.py and add gradient accumulation loop as above.
6. Run with TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE="xla_graph_executor=6,pjrt_computation_client=5" to see aliasing debugging logs:

XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 XLA_SAVE_TENSORS_FMT="hlo" XLA_SAVE_TENSORS_FILE="/tmp/save1.hlo"   TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE="xla_graph_executor=6,pjrt_computation_client=5" python test_train_mp_mnist.py |& tee log

Expected behavior

Expect gradients to be aliased

Environment

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

Additional context

[JackCaoG]

so I tired to test it on nightly, with

diff --git a/examples/train_resnet_base.py b/examples/train_resnet_base.py
index b66780d5c..acd00edff 100644
--- a/examples/train_resnet_base.py
+++ b/examples/train_resnet_base.py
@@ -11,6 +11,8 @@ import torch_xla
 import torchvision
 import torch.optim as optim
 import torch.nn as nn
+import torch_xla.debug.metrics as met
+
 
 
 class TrainResNetBase():
@@ -29,7 +31,7 @@ class TrainResNetBase():
         xr.world_size())
 
     self.device = torch_xla.device()
-    self.train_device_loader = pl.MpDeviceLoader(train_loader, self.device)
+    self.train_device_loader = pl.MpDeviceLoader(train_loader, self.device, batches_per_execution=4)
     self.model = torchvision.models.resnet50().to(self.device)
     self.optimizer = optim.SGD(self.model.parameters(), weight_decay=1e-4)
     self.loss_fn = nn.CrossEntropyLoss()
@@ -51,6 +53,8 @@ class TrainResNetBase():
       loss.backward()
       self.run_optimizer()
       tracker.add(self.batch_size)
+      print('alias count = ')
+      print(met.metric_data("InputOutputAliasCount"))
       if step % 10 == 0:
         xm.add_step_closure(
             self._train_update, args=(step, loss, tracker, epoch))

we can do graident accumulation for every 4 steps for resnet50. When run with

PT_XLA_DEBUG_LEVEL=1 python examples/train_resnet_base.py 

I see

Compilation Analysis: ================================================================================
Compilation Analysis: Compilation Cause
Compilation Analysis:   mark_step in parallel loader at step end
Compilation Analysis: Graph Info: 
Compilation Analysis:   Graph Hash: e58a0059c5ae1ac0035a1d6fcefbc377
Compilation Analysis:   Number of Graph Inputs: 272
Compilation Analysis:   Number of Graph Outputs: 486
Compilation Analysis: Python Frame Triggered Execution: 
Compilation Analysis:   mark_step (/workspaces/dk3/pytorch/xla/torch_xla/core/xla_model.py:1055)
Compilation Analysis:   next (/workspaces/dk3/pytorch/xla/torch_xla/distributed/parallel_loader.py:44)
Compilation Analysis:   __next__ (/workspaces/dk3/pytorch/xla/torch_xla/distributed/parallel_loader.py:32)
Compilation Analysis:   train_loop_fn (/workspaces/dk3/pytorch/xla/examples/train_resnet_base.py:49)
Compilation Analysis:   start_training (/workspaces/dk3/pytorch/xla/examples/train_resnet_base.py:67)
Compilation Analysis:   <module> (/workspaces/dk3/pytorch/xla/examples/train_resnet_base.py:75)
Compilation Analysis: --------------------------------------------------------------------------------
Compilation Analysis: ================================================================================

Post Compilation Analysis: ================================================================================
Post Compilation Analysis: Graph input size: 0.167664 GB
Post Compilation Analysis: Graph output size: 0.199731 GB
Post Compilation Analysis: Aliased Input size: 0.095884 GB
Post Compilation Analysis: Intermediate tensor size: 8.436128 GB
Post Compilation Analysis: Compiled program size: 0.194783 GB
Post Compilation Analysis: --------------------------------------------------------------------------------
Post Compilation Analysis: ================================================================================
epoch: 1, step: 0, loss: 6.9063191413879395, rate: 8.059381469430795
alias count = 
(1, 267.0, ((1717453831.561344, 267.0),))

Note that 0.095884 GB of input is alised, out of 0.167664 GB. and 267/272 of inputs are being aliased. now if I run it again with

XLA_DISABLE_FUNCTIONALIZATION=1 PT_XLA_DEBUG_LEVEL=1 python examples/train_resnet_base.py

I see exactly the same thing.

Post Compilation Analysis: ================================================================================
Post Compilation Analysis: Graph input size: 0.167664 GB
Post Compilation Analysis: Graph output size: 0.192098 GB
Post Compilation Analysis: Aliased Input size: 0.095884 GB
Post Compilation Analysis: Intermediate tensor size: 8.455472 GB
Post Compilation Analysis: Compiled program size: 0.194823 GB
Post Compilation Analysis: --------------------------------------------------------------------------------
Post Compilation Analysis: ================================================================================
epoch: 1, step: 0, loss: 6.912147045135498, rate: 8.488426770395343
alias count = 
(1, 267.0, ((1717454064.4393814, 267.0),))

This suggest on nightly XLA_DISABLE_FUNCTIONALIZATION does not impact aliasing. It is hard to fix the 2.1 branch at this point. I would suggest just run with XLA_DISABLE_FUNCTIONALIZATION=1 until you guys upgraded to 2.3 release.

[JackCaoG]

Ok here is how I will debug this problem. Let's first understand when and how do we do buffer aliasing in 2.1 branch.

xla/torch_xla/csrc/xla_graph_executor.cpp

Lines 1370 to 1371 in b7ea06e

	input_output_alias_pair =
	BuildInputOutputAliases(tensors, coll.indices, &lowering_ctx);

Build input output alias happens during the compilation, so it is a compile time decision.
https://github.com/pytorch/xla/blob/b7ea06efbf0bb1ae1f094c07db7387470cd787ee/torch_xla/csrc/xla_graph_executor.cpp#L1234C19-L1236
It takes tensor that represent outputs (tensors + indices are all XLATensors that will get a new output value after current execution) and lowering context as input.

we first find the tensor ID of all output tensors

xla/torch_xla/csrc/xla_graph_executor.cpp

Lines 1242 to 1246 in b7ea06e

	for (size_t i = 0; i < indices.size(); ++i) {
	size_t tensor_index = indices[i];
	int64_t tensor_id = tensors[tensor_index]->data()->alias_id;
	output_tensor_id_map[tensor_id] = i;
	}

and then we loop through all of the input datas(from the lowering context)

xla/torch_xla/csrc/xla_graph_executor.cpp

Lines 1247 to 1249 in b7ea06e

	const auto& parameters_data = lowering_ctx->GetParametersData();
	std::vector<ssize_t> alias_map(indices.size(), -1);
	for (size_t i = 0; i < parameters_data.size(); ++i) {

and see if there is any input data share the same tensorID as the output. If they share the same input ID it means that some in place operation happened and XLATensor no longer attached to this input data so this input data can be aliased to the output with the same shape.

Note we make a couple assumption here

1. There is only one XLATensor associated with the one XLAData which is usually the case
2. XLAData's ID correctly reflect the owner of the XLATensor

In order to debug issue above, I think what you want to do is

1. dump all of output tensor's TensorID
2. dump all of input data's TensorID

and see why don't they match

[jeffhataws]

Just want to document some findings using the original MLP test using just CPU, printing met.metric_data("InputOutputAliasCount"):

pt2.1 with functionalization on CPU:
(4, 42.0, ((1719204665.609187, 16.0), (1719204667.7811365, 2.0), (1719204668.486205, 12.0), (1719204668.6342065, 12.0))) 

pt2.1 without functionalization on CPU:
(4, 58.0, ((1719204732.9488149, 16.0), (1719204735.0688317, 6.0), (1719204735.781659, 12.0), (1719204735.9386797, 24.0)))

pt2.3 with functionalization on CPU:
(4, 42.0, ((1719204826.4645836, 16.0), (1719204829.2127182, 2.0), (1719204829.8989058, 12.0), (1719204830.0465565, 12.0)))

pt2.3 without functionalization on CPU:
(4, 58.0, ((1719204809.890676, 16.0), (1719204812.663207, 6.0), (1719204813.3367858, 12.0), (1719204813.4833145, 24.0)))  

pt2.4 with functionalization on CPU:
(4, 54.0, ((1719207116.7393954, 16.0), (1719207119.7452862, 2.0), (1719207120.3934603, 12.0), (1719207120.5731502, 24.0)))

pt2.4 without functionalization on CPU:
(4, 58.0, ((1719207140.930515, 16.0), (1719207143.827449, 6.0), (1719207144.4554996, 12.0), (1719207144.6054273, 24.0)))

[jeffhataws]

Minimal reproduction with only 1 linear layer and only gradient accumulation:

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_xla.core.xla_model as xm
import torch_xla.debug.metrics as met
import numpy as np
import copy

dev = xm.xla_device()

train_x_sample = torch.rand((1, 28 * 28))
train_label_sample = torch.tensor([5])

class MLP(nn.Module):
  def __init__(self, input_size = 28 * 28, output_size = 10):
      super(MLP, self).__init__()
      self.fc1 = nn.Linear(input_size, output_size, bias=False)

  def forward(self, x):
      x = self.fc1(x)
      return F.log_softmax(x, dim=1)

c_model = MLP().to('cpu')
t_model = copy.deepcopy(c_model).to(dev)
loss_fn = nn.NLLLoss()

def try_grad_accum(model, device, train_x, train_label, accum_steps):
    train_x = train_x.to(device)
    train_label = train_label.to(device)
    model.zero_grad()
    for i in range(accum_steps):
        output = model(train_x)
        t_loss = loss_fn(output, train_label)
        t_loss.backward()
        xm.mark_step()
    return [p.grad.to('cpu').numpy() for p in model.parameters()]

def test_grad_accum():
    t_model.train()
    c_model.train()
    accum_steps = 4
    c_grads_5 = try_grad_accum(c_model, 'cpu', train_x_sample, train_label_sample, accum_steps)
    t_grads_5 = try_grad_accum(t_model, dev, train_x_sample, train_label_sample, accum_steps)
    np.testing.assert_allclose(t_grads_5, c_grads_5, rtol=3e-2, atol=1e-3)
    print("XLA_DISABLE_FUNCTIONALIZATION:", os.environ["XLA_DISABLE_FUNCTIONALIZATION"],
          " InputOutputAliasCount:", met.metric_data("InputOutputAliasCount"))

if __name__ == '__main__':
    test_grad_accum()

Disable functionalization (see aliasing of gradient buffer):

XLA_FLAGS=--xla_dump_to="./dump_nofunc" XLA_DISABLE_FUNCTIONALIZATION=1 python gradacc_test.py

XLA_DISABLE_FUNCTIONALIZATION: 1  InputOutputAliasCount: (2, 1.0, ((1719295462.695567, 0.0), (1719295462.758401, 1.0)))

dump_nofunc/module_0001.SyncTensorsGraph.115.cpu_after_optimizations.txt also shows aliasing:

HloModule SyncTensorsGraph.115, is_scheduled=true, input_output_alias={ {0}: (3, {}, may-alias) }, entry_computation_layout={(f32[10,784]{1,0}, f32[1,784]{1,0}, s64[1]{0}, f32[10,784]{1,0})->(f32[10,784]{1,0}, f32[1,10]{1,0}, f32[])}

Enable functionalization (see no aliasing):

XLA_FLAGS=--xla_dump_to="./dump_func" XLA_DISABLE_FUNCTIONALIZATION=0 python gradacc_test.py

XLA_DISABLE_FUNCTIONALIZATION: 0  InputOutputAliasCount: (2, 0.0, ((1719295465.6997023, 0.0), (1719295465.7658002, 0.0)))

dump_func/module_0001.SyncTensorsGraph.117.cpu_after_optimizations.txt shows no aliasing:

HloModule SyncTensorsGraph.117, is_scheduled=true, entry_computation_layout={(f32[10,784]{1,0}, f32[1,784]{1,0}, s64[1]{0}, f32[10,784]{1,0})->(f32[1,10]{1,0}, f32[], f32[784,10]{1,0}, f32[10,784]{1,0})}

[jeffhataws]

Using TOT, I modified torch_xla/csrc/xla_graph_executor.cpp to dump some info:

diff --git a/torch_xla/csrc/xla_graph_executor.cpp b/torch_xla/csrc/xla_graph_executor.cpp
index 74c3270a9..d31924fb4 100644
--- a/torch_xla/csrc/xla_graph_executor.cpp
+++ b/torch_xla/csrc/xla_graph_executor.cpp
@@ -1264,6 +1264,7 @@ std::vector<size_t> XLAGraphExecutor::SetBufferDonors(
     size_t tensor_index = indices[i];
     int64_t tensor_id = tensors[tensor_index]->data()->alias_id;
     output_tensor_id_map[tensor_id] = i;
+    TF_VLOG(3) << "SetBufferDonors indices_index: " << i << " tensor_index: " << tensor_index << " tensor_id: " << tensor_id;
   }
   const auto& parameters_data = lowering_ctx->GetParametersData();
   std::vector<ssize_t> alias_map(indices.size(), -1);
@@ -1271,8 +1272,10 @@ std::vector<size_t> XLAGraphExecutor::SetBufferDonors(
     auto* data_info =
         static_cast<torch::lazy::LazyGraphExecutor::DeviceDataInfo*>(
             parameters_data[i]->info());
+    TF_VLOG(3) << "SetBufferDonors param index: " << i << " data_info: " << data_info;
     if (data_info != nullptr && !data_info->read_only) {
       auto it = output_tensor_id_map.find(data_info->tensor_id);
+      TF_VLOG(3) << "SetBufferDonors data_info->tensor_id: " << data_info->tensor_id;
       // Parameter buffer's TensorId in output_tensor_id_map means
       // this buffer is not needed after execution since XLATensor will get a
       // new buffer.
@@ -1284,6 +1287,7 @@ std::vector<size_t> XLAGraphExecutor::SetBufferDonors(
     }
   }
   TORCH_LAZY_VALUE_METRIC("InputOutputAliasCount", buffer_donor_indexs.size());
+  TF_VLOG(3) << "SetBufferDonors buffer_donor_indexs: " << buffer_donor_indexs;
   return buffer_donor_indexs;
 }

Disable functionalization (timestamps removed to improve readabiility):

torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 0 tensor_index: 3 tensor_id: 6
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 1 tensor_index: 4 tensor_id: 8
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 2 tensor_index: 5 tensor_id: 14
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 0 data_info: 0x5583fe936890
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 1
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 1 data_info: 0x5583fe968640
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 2
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 2 data_info: 0x5583fe902550
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 3
torch_xla/csrc/xla_graph_executor.cpp:1293] SetBufferDonors buffer_donor_indexs:
torch_xla/csrc/xla_graph_executor.cpp:1421] Compiling IR graph hash a7fe2231121157aab2b76c36f9085caf on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash a7fe2231121157aab2b76c36f9085caf on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 0 tensor_index: 3 tensor_id: 14
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 1 tensor_index: 4 tensor_id: 17
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 2 tensor_index: 5 tensor_id: 19
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 0 data_info: 0x5583fe936890
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 1
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 1 data_info: 0x5583fe968640
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 2
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 2 data_info: 0x5583fe903160
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 3
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 3 data_info: 0x7f7910019030
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 14
torch_xla/csrc/xla_graph_executor.cpp:1293] SetBufferDonors buffer_donor_indexs: 3
torch_xla/csrc/xla_graph_executor.cpp:1421] Compiling IR graph hash 8c5d2f3da9d51ab21d4f711702e98d0c on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 8c5d2f3da9d51ab21d4f711702e98d0c on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 8c5d2f3da9d51ab21d4f711702e98d0c on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 8c5d2f3da9d51ab21d4f711702e98d0c on device CPU:0 done!
XLA_DISABLE_FUNCTIONALIZATION: 1  InputOutputAliasCount: (2, 1.0, ((1719438778.8729675, 0.0), (1719438778.9325185, 1.0)))

With functionalization:

torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 0 tensor_index: 3 tensor_id: 6
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 1 tensor_index: 4 tensor_id: 8
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 2 tensor_index: 5 tensor_id: 13
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 3 tensor_index: 6 tensor_id: 14
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 0 data_info: 0x5556f7ef6d60
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 1
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 1 data_info: 0x5556f7f5ceb0
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 2
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 2 data_info: 0x5556f7ef5910
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 3
torch_xla/csrc/xla_graph_executor.cpp:1293] SetBufferDonors buffer_donor_indexs:
torch_xla/csrc/xla_graph_executor.cpp:1421] Compiling IR graph hash 86eb0fe9e6afb2f4cec87b9250f18010 on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 86eb0fe9e6afb2f4cec87b9250f18010 on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 0 tensor_index: 3 tensor_id: 17
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 1 tensor_index: 4 tensor_id: 19
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 2 tensor_index: 5 tensor_id: 28
torch_xla/csrc/xla_graph_executor.cpp:1269] SetBufferDonors indices_index: 3 tensor_index: 6 tensor_id: 29
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 0 data_info: 0x5556f7ef6d60
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 1
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 1 data_info: 0x5556f7f4a180
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 2
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 2 data_info: 0x5556f7f577e0
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 3
torch_xla/csrc/xla_graph_executor.cpp:1277] SetBufferDonors param index: 3 data_info: 0x7fb9c4018300
torch_xla/csrc/xla_graph_executor.cpp:1280] SetBufferDonors data_info->tensor_id: 14
torch_xla/csrc/xla_graph_executor.cpp:1293] SetBufferDonors buffer_donor_indexs:
torch_xla/csrc/xla_graph_executor.cpp:1421] Compiling IR graph hash 578c2535b5823addcb078d3b876947e1 on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 578c2535b5823addcb078d3b876947e1 on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 578c2535b5823addcb078d3b876947e1 on device CPU:0 done!
torch_xla/csrc/xla_graph_executor.cpp:1132] Executing IR graph hash 578c2535b5823addcb078d3b876947e1 on device CPU:0 done!
XLA_DISABLE_FUNCTIONALIZATION: 0  InputOutputAliasCount: (2, 0.0, ((1719438794.6676075, 0.0), (1719438794.7328753, 0.0)))

[jeffhataws]

In functionalization graph, there's one additional output of shape (768, 10) that's not there in no-functionalization case.

[JackCaoG]

Ok I was able to repo and I had a rought idea about what's going on. This bug seems to only happen with the graident accumulation.

so we first have the model.fc1.weight.grad after the first step and in the second step PyTorch will use += to accumulate the grad. The tensor ID for model.fc1.weight.grad was 14 before the second step, and during the backward, PyTorch/XLA recevied a call to XLANativeFunctions::_propagate_xla_data which tell us some in place operation happens bettwen tensor ID 14 and and tensor ID 26. PyTorch/XLA learned to link these two tensors together. However after the backward is finished, the tensor ID of the new model.fc1.weight.grad becomes 29. In this case PyTorch/XLA can't tell that tensor ID 14 and tensor ID 29 is the same tensor hence it will not aliased.

@bdhirsh can you point me to where the graident accumulation happens in Pytorch? It happens in C++ so it is a bit hard for me to trace the exact location without doing debug build for both PyTorch and Pytorch/XLA.

My theory is that functionization did something wierd for the graident accumulation which confuse PyTorch/XLA so we can't really tell this is an inplace call.

[JackCaoG]

@jeffhataws what I did was adding

print(torch_xla._XLAC._get_xla_tensor_debug_info(model.fc1.weight.grad))

after the mark_step and add following debug prints

diff --git a/torch_xla/csrc/aten_xla_type.cpp b/torch_xla/csrc/aten_xla_type.cpp
index 5e233d051..e222231f3 100644
--- a/torch_xla/csrc/aten_xla_type.cpp
+++ b/torch_xla/csrc/aten_xla_type.cpp
@@ -2707,6 +2707,12 @@ void XLANativeFunctions::_propagate_xla_data(const at::Tensor& input,
   // 1) Aid XLA's InputOutputAlias.
   auto input_tensor = bridge::GetXlaTensor(input);
   auto output_tensor = bridge::GetXlaTensor(output);
+  std::cout << "output tensor's alias ID: " << output_tensor->data()->alias_id
+            << "\n";
+  std::cout << "input tensor's tensor ID: " << input_tensor->GetUniqueId()
+            << "\n";
+  std::cout << "input tensor's alias ID: " << input_tensor->data()->alias_id
+            << "\n";
   if (input_tensor->CurrentDataHandle() != nullptr ||
       (input_tensor->CurrentIrValue().node != nullptr &&
        torch_xla::DeviceData::Cast(
@@ -2722,7 +2728,7 @@ void XLANativeFunctions::_propagate_xla_data(const at::Tensor& input,
     xm.mark_step()
     // x.tensor_id =3, x.alias_id should be 2 since input tensor id will be 2
     // for this graph
-    x *= 1 of 1
+    x *= 1
     */
     output_tensor->data()->alias_id = input_tensor->GetUniqueId();
   } else {
diff --git a/torch_xla/csrc/xla_graph_executor.cpp b/torch_xla/csrc/xla_graph_executor.cpp
index 0cbb21c44..bd8a9af06 100644
--- a/torch_xla/csrc/xla_graph_executor.cpp
+++ b/torch_xla/csrc/xla_graph_executor.cpp
@@ -1275,6 +1275,8 @@ std::vector<size_t> XLAGraphExecutor::SetBufferDonors(
   for (size_t i = 0; i < indices.size(); ++i) {
     size_t tensor_index = indices[i];
     int64_t tensor_id = tensors[tensor_index]->data()->alias_id;
+    std::cout << "adding alias_id " << tensor_id << " to the map"
+              << "\n";
     output_tensor_id_map[tensor_id] = i;
   }
   const auto& parameters_data = lowering_ctx->GetParametersData();
@@ -1283,6 +1285,7 @@ std::vector<size_t> XLAGraphExecutor::SetBufferDonors(
     auto* data_info =
         static_cast<torch::lazy::LazyGraphExecutor::DeviceDataInfo*>(
             parameters_data[i]->info());
+    std::cout << "parameter tensor id = " << data_info->tensor_id << "\n";
     if (data_info != nullptr && !data_info->read_only) {
       auto it = output_tensor_id_map.find(data_info->tensor_id);
       // Parameter buffer's TensorId in output_tensor_id_map means

[bdhirsh]

Hmm @JackCaoG the grad accumulation impl is here: https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/functions/accumulate_grad.h#L116

[jeffhataws]

One additional datapoint is that when I increase the gradient accumulation count, I see that the input tensor ID keeps changing for each gradient accumulation step when functionalization is on. Is that expected?

xla_graph_executor.cpp:1267] SetBufferDonors: adding alias_id 17 to the map to indices_index 0
xla_graph_executor.cpp:1268] SetBufferDonors indices_index: 0 tensor_index: 3 tensor_id(alias_id): 17
xla_graph_executor.cpp:1267] SetBufferDonors: adding alias_id 19 to the map to indices_index 1
xla_graph_executor.cpp:1268] SetBufferDonors indices_index: 1 tensor_index: 4 tensor_id(alias_id): 19    
xla_graph_executor.cpp:1267] SetBufferDonors: adding alias_id 28 to the map to indices_index 2             
xla_graph_executor.cpp:1268] SetBufferDonors indices_index: 2 tensor_index: 5 tensor_id(alias_id): 28
xla_graph_executor.cpp:1267] SetBufferDonors: adding alias_id 29 to the map to indices_index 3
xla_graph_executor.cpp:1268] SetBufferDonors indices_index: 3 tensor_index: 6 tensor_id(alias_id): 29
xla_graph_executor.cpp:1276] SetBufferDonors param index: 0 data_info->tensor_id: 1
xla_graph_executor.cpp:1279] SetBufferDonors data_info->tensor_id: 1
xla_graph_executor.cpp:1276] SetBufferDonors param index: 1 data_info->tensor_id: 2
xla_graph_executor.cpp:1279] SetBufferDonors data_info->tensor_id: 2
xla_graph_executor.cpp:1276] SetBufferDonors param index: 2 data_info->tensor_id: 3
xla_graph_executor.cpp:1279] SetBufferDonors data_info->tensor_id: 3
xla_graph_executor.cpp:1276] SetBufferDonors param index: 3 data_info->tensor_id: 14
xla_graph_executor.cpp:1279] SetBufferDonors data_info->tensor_id: 14
xla_graph_executor.cpp:1291] SetBufferDonors buffer_donor_indexs: 
xla_graph_executor.cpp:1412] Compiling IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 ...
xla_graph_executor.cpp:1419] Compiling IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 done!
xla_graph_executor.cpp:1119] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 ...
xla_graph_executor.cpp:1130] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 done!
aten_xla_type.cpp:2714] output tensor's alias ID: 41 shape: f32[10,784]
aten_xla_type.cpp:2715] input tensor's tensor ID: 29
aten_xla_type.cpp:2716] input tensor's alias ID: 29 shape: f32[10,784]
aten_xla_type.cpp:2735] assigning output tensor's alias ID to input tensor's (unique) ID
xla_graph_executor.cpp:1119] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 ...
xla_graph_executor.cpp:1130] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 done!
aten_xla_type.cpp:2714] output tensor's alias ID: 56 shape: f32[10,784]
aten_xla_type.cpp:2715] input tensor's tensor ID: 44
aten_xla_type.cpp:2716] input tensor's alias ID: 44 shape: f32[10,784]
aten_xla_type.cpp:2735] assigning output tensor's alias ID to input tensor's (unique) ID
xla_graph_executor.cpp:1119] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 ...
xla_graph_executor.cpp:1130] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 done!
aten_xla_type.cpp:2714] output tensor's alias ID: 71 shape: f32[10,784]
aten_xla_type.cpp:2715] input tensor's tensor ID: 59
aten_xla_type.cpp:2716] input tensor's alias ID: 59 shape: f32[10,784]
aten_xla_type.cpp:2735] assigning output tensor's alias ID to input tensor's (unique) ID
xla_graph_executor.cpp:1119] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 ...
xla_graph_executor.cpp:1130] Executing IR graph hash 417805772c05e5e266b5983e46792f74 on device CPU:0 done!
aten_xla_type.cpp:2714] output tensor's alias ID: 86 shape: f32[10,784]
aten_xla_type.cpp:2715] input tensor's tensor ID: 74
aten_xla_type.cpp:2716] input tensor's alias ID: 74 shape: f32[10,784]
aten_xla_type.cpp:2735] assigning output tensor's alias ID to input tensor's (unique) ID

[JackCaoG]

tensor ID keep changing is expected, with functionalization there is no real in place ops. We can still alias the underlying buffer as long as we know they are supposely the same tensor(through _propagate_xla_data ).

[jeffhataws]

Make sense. I think for the second graph, the input tensor_id 14 should be aliased to tensor_id 3, instead of aliased to itself (as indicated by the map in SetBufferDonors: "adding alias_id 14 to the map to indices_index 3"). This way, tensor_id 3 would keep propagating.

[JackCaoG]

I found where tensor ID 27 -29 were from

131│ bool FunctionalStorageImpl::apply_updates() {
132│   // N.B:none of the tensors used in this function should be FunctionalTensorWrappers at this point.
133│   // The only reason we currently need the TLS exclude guard here is because of functorch's DynamicLayer stack.
134│   // It adds the Functionalize key into TLS before redispatching to the functionalization kernels,
135│   // which means that we need to explicitly exclude it here before doing any other work underneath the pass.
136│   at::AutoDispatchSkipFunctionalize guard;
137│   bool any_updates = !updates_.empty();
138│   for (auto& update_data: updates_) {
139├>    base_ = apply_update(update_data, base_);
140│   }
141│   updates_.clear();
142│   return any_updates;
143│ }

will trigger t_copy in pytorch/xla

and it is from

 2252│   at::functionalization::impl::replace_(self, tmp_output);
 2253│   at::functionalization::impl::commit_update(self);
 2254├>  at::functionalization::impl::sync(self);
 2255│     return self;

which is part of

 2204│     at::Tensor & add__Tensor(c10::DispatchKeySet dispatchKeySet, at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {

in pytorch/build/aten/src/ATen/RegisterFunctionalization_0.cpp .

Here is what happened. Tensor with ID 26 was created with ir add, this is due to the add_ in the accumulated grad. _propagate_xla_data was correctly called and pytorch/xla knows that 26 and 14 was the same tensor. However after the add is finished, at::functionalization::impl::sync triggered a couple permute op and inplace replace the base_ with new tensors. There is no way PyTorch/XLA knows that new permute IR is the same tensor as the add tensor.

@bdhirsh can we also call the _propagate_xla_data for base_ before and after apply_update?

[jeffhataws]

Thanks @JackCaoG for the detailed root-cause. @bdhirsh do you think this can be fix in torch 2.4?