Skip to content

Commit

Permalink
Working on transpose/training tests.
Browse files Browse the repository at this point in the history
Updated training test.
  • Loading branch information
vladimirjovanovicTT committed Aug 22, 2024
1 parent fffeed3 commit 502a51d
Show file tree
Hide file tree
Showing 6 changed files with 44 additions and 19 deletions.
11 changes: 8 additions & 3 deletions pybuda/csrc/buda_passes.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -148,8 +148,11 @@ void run_optimization_graph_passes(graphlib::Graph *graph)
if (not skip_erase_redundant) {
if (not attempt_update)
attempt_update = passes::erase_consecutive_reshape(graph, true);
if (not attempt_update)
attempt_update = passes::fuse_tm_sequences(graph);

// TODO: Figure out if this is needed. (Issue #152)
// if (not attempt_update)
// attempt_update = passes::fuse_tm_sequences(graph);

passes::bypass_nop_tms(graph);
}
}
Expand All @@ -167,7 +170,9 @@ void run_optimization_graph_passes(graphlib::Graph *graph)

passes::move_select_after_matmul_optional(graph);

passes::fuse_tm_sequences(graph);
// Issue #152
// passes::fuse_tm_sequences(graph);

reportify::dump_graph(graph->name(), "post_erase_inverse_ops", graph);
}

Expand Down
1 change: 1 addition & 0 deletions pybuda/csrc/passes/lower_to_mlir.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -504,6 +504,7 @@ class MLIRGenerator
lowering_handler_map["softmax"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SoftmaxOp>;
lowering_handler_map["reduce_sum"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SumOp>;
lowering_handler_map["reduce_avg"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::MeanOp>;
lowering_handler_map["transpose"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::TransposeOp>;
}
};
}
Expand Down
2 changes: 1 addition & 1 deletion pybuda/pybuda/op/tm.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ def Transpose(
if dim0 > dim1:
dim0, dim1 = dim1, dim0

return op("transpose", name, operandA, dim0=dim0, dim1=dim1, z_dim_slice=z_dim_slice).get_tensor(out_df=pytorch_dtype_to_buda_dataformat(out_dtype))
return op("transpose", name, operandA, attrs=(dim0, dim1, z_dim_slice), dim0=dim0, dim1=dim1).get_tensor(out_df=pytorch_dtype_to_buda_dataformat(out_dtype))

def Reshape(
name: str,
Expand Down
19 changes: 19 additions & 0 deletions pybuda/test/mlir/test_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,25 @@ def forward(self, a, b):
co_out = [co.to("cpu") for co in co_out]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]

def test_transpose():
class Transpose(nn.Module):
def __init__(self):
super().__init__()

def forward(self, a, b):
c = a + b
return torch.transpose(c, 1, 2)

inputs = [torch.rand(1, 32, 64), torch.rand(1, 32, 64)]

framework_model = Transpose()
fw_out = framework_model(*inputs)

compiled_model = pybuda.compile(framework_model, sample_inputs=inputs)
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]

def test_subtract():
class Subtract(nn.Module):
Expand Down
28 changes: 14 additions & 14 deletions pybuda/test/mlir/test_training.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,51 +7,52 @@

import pybuda
import pybuda.config
from pybuda.op.eval.common import compare_with_golden_pcc

def test_torch_training():
class MultParam(nn.Module):
def __init__(self):
super().__init__()
self.p = nn.Parameter(torch.rand(1, 1024))
self.p = nn.Parameter(torch.ones(1024, 1024))
nn.init.xavier_uniform_(self.p)

def forward(self, x1):
return torch.multiply(x1, self.p)
return torch.matmul(x1, self.p)

model = MultParam()
shape = (1, 1024)
shape = (8, 1024)
inputs = torch.rand(shape)

# Fake targets
target = torch.zeros(shape)

loss_fn = torch.nn.L1Loss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

tt_model = pybuda.compile(model, sample_inputs=[torch.rand(shape)], loss=loss_fn, optimizer=optimizer)

num_epochs = 100
num_epochs = 20

for epoch in range(num_epochs):

print(f"parameter value: {model.p.data}")
#print(f"parameter value: {model.p.data}")
golden = model(inputs)
output = tt_model(inputs)

if not torch.allclose(output[0], golden, rtol=1e-1):
raise ValueError("Output does not match the golden output")
#print(f"golden = {golden}, output = {output[0]}")
oputput = [co.to("cpu") for co in output]
assert compare_with_golden_pcc(golden=golden, calculated=output[0], pcc=0.99)

optimizer.zero_grad()

loss = loss_fn(output[0], target)
loss.backward()

print(f"epoch: {epoch} loss: {loss}")
print(f"output.grad: {output[0].grad}")
#print(f"output.grad: {output[0].grad}")

loss_grad = output[0].grad
assert loss_grad is not None

print(f"loss grad: {loss_grad}")
#print(f"loss grad: {loss_grad}")
grad = tt_model.backward(loss_grad)

# HACK to run the optimizer step
Expand All @@ -60,4 +61,3 @@ def forward(self, x1):
model.p.grad = grad[0]

optimizer.step()

2 changes: 1 addition & 1 deletion third_party/tt-mlir
Submodule tt-mlir updated 121 files

0 comments on commit 502a51d

Please sign in to comment.