Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Track All Changes from upstream #45

Open
wants to merge 52 commits into
base: origin
Choose a base branch
from
Open
Changes from 1 commit
Commits
Show all changes
52 commits
Select commit Hold shift + click to select a range
23e38ff
Backend Cleanroom port (#13)
MarisaKirisame Mar 6, 2020
5ffafc2
Hook up the overload with a passing test. (#15)
MarisaKirisame Mar 16, 2020
6464752
Adding some overload (#16)
MarisaKirisame Mar 20, 2020
d09fb75
make resnet pass again, adding more logs for logging simulator (#17)
MarisaKirisame Mar 26, 2020
6e6c4e1
Logging everything for resnet (#18)
MarisaKirisame Apr 8, 2020
dab0aa8
Fix logger (special handling for small constants) (#19)
MarisaKirisame Apr 15, 2020
0680ddb
TreeLSTM overloads (#21)
MarisaKirisame Apr 17, 2020
3008ba5
Overloads for mean and mean.dim (needed for densenet)
slyubomirsky Apr 18, 2020
97c7265
Merge pull request #22 from slyubomirsky/densenet-overload
MarisaKirisame Apr 18, 2020
b2cc5f4
Add overloads for U-Net
slyubomirsky Apr 30, 2020
4324672
Merge pull request #24 from slyubomirsky/unet-overloads
MarisaKirisame Apr 30, 2020
5c23115
Add ability to annotate log and make new log file (#25)
MarisaKirisame May 5, 2020
690883f
Restore overloads needed for LSTM and GRU
slyubomirsky May 6, 2020
5d8cfd1
Merge pull request #27 from slyubomirsky/lstm-overloads
MarisaKirisame May 7, 2020
d5f5e86
Implemented unrolled_gan (#23)
MarisaKirisame May 7, 2020
d8a4aea
Overload bitwise and for ParityACT
slyubomirsky May 7, 2020
97fb92d
Add various overloads for ACT model
slyubomirsky May 13, 2020
2f719fb
Merge pull request #30 from slyubomirsky/act-overloads
MarisaKirisame May 13, 2020
74ce175
Log aliases of operator outputs
MarisaKirisame May 13, 2020
3545dfb
Overloads and changes needed for transformer
slyubomirsky May 15, 2020
ae01104
Merge pull request #33 from slyubomirsky/equal-overload
MarisaKirisame May 15, 2020
79578bb
Overloads for topk functions
slyubomirsky May 16, 2020
692a346
Merge pull request #34 from slyubomirsky/topk-overload
MarisaKirisame May 17, 2020
e144342
Add overloads for deepspeech
slyubomirsky May 19, 2020
5476fcc
More overloads for ACT
slyubomirsky May 20, 2020
224e6d9
Merge pull request #35 from slyubomirsky/deepspeech-support
MarisaKirisame May 21, 2020
0b770f3
DTR implementation first pass (#37)
MarisaKirisame May 23, 2020
73f681c
refactor - remove stitch
MarisaKirisame May 23, 2020
ec0fe99
save
MarisaKirisame May 23, 2020
0118183
Merge pull request #38 from uwsampl/refactor
MarisaKirisame May 24, 2020
41f694b
save
MarisaKirisame May 24, 2020
f17f292
[ impl ] overload diag & mv
AD1024 May 28, 2020
e4c3ab6
Merge pull request #40 from uwsampl/diag-mv-overload
MarisaKirisame May 28, 2020
7c5f36d
refactor - remove stitch
MarisaKirisame May 23, 2020
c9ef097
save
MarisaKirisame May 31, 2020
0a08ff7
save
MarisaKirisame Jun 1, 2020
fe395f4
save
MarisaKirisame Jun 1, 2020
7efe202
More opts
jroesch Jun 2, 2020
ed31496
Try removing pop-ing from loop
jroesch Jun 2, 2020
a9878a5
save
MarisaKirisame Jun 2, 2020
f92f490
Merge branch 'search' of github.com:uwsampl/pytorch into search
MarisaKirisame Jun 2, 2020
984ec10
save
MarisaKirisame Jun 2, 2020
4607be4
save
MarisaKirisame Jun 2, 2020
8ad4f93
save
MarisaKirisame Jun 2, 2020
74dd61c
save
MarisaKirisame Jun 3, 2020
f8c8f32
mike try this
MarisaKirisame Jun 3, 2020
f2f095c
mike try this
MarisaKirisame Jun 3, 2020
7c04070
save
MarisaKirisame Jun 3, 2020
3cb823b
undo system optimization for now
MarisaKirisame Jun 3, 2020
f6312c9
save
MarisaKirisame Jun 4, 2020
ee8f2a9
save
MarisaKirisame Jun 4, 2020
4540277
save
MarisaKirisame Jun 23, 2020
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Overloads and changes needed for transformer
  • Loading branch information
slyubomirsky committed May 15, 2020
commit 3545dfba0f420ed445da2573478b05d3a1da1e94
52 changes: 52 additions & 0 deletions aten/src/ATen/native/Checkpoint.cpp
Original file line number Diff line number Diff line change
@@ -1203,6 +1203,58 @@ Tensor& checkpoint_index_put_(Tensor& self, ArrayRef<Tensor> indices, const Tens
return self;
}

Tensor checkpoint_bmm(const Tensor& self, const Tensor& mat2) {
rematerialize_function_t rt =
[=](const Tensors& vec) -> Tensors {
return {at::bmm(vec.at(0), vec.at(1))};
};
return CheckpointTensorImpl::make("bmm", rt, {self, mat2})[0];
}

Tensor checkpoint__softmax(const Tensor& self, long dim, bool half_to_float) {
rematerialize_function_t rt =
[=](const Tensors& vec) -> Tensors {
return {at::_softmax(vec.at(0), dim, half_to_float)};
};
return CheckpointTensorImpl::make("_softmax", rt, {self})[0];
}

Tensor checkpoint__softmax_backward_data(const Tensor& grad_output, const Tensor& output, long dim, const Tensor& self) {
rematerialize_function_t rt =
[=](const Tensors& vec) -> Tensors {
return {at::_softmax_backward_data(vec.at(0), vec.at(1), dim, vec.at(2))};
};
return CheckpointTensorImpl::make("_softmax_backward_data", rt, {grad_output, output, self})[0];
}

std::tuple<Tensor, Tensor, Tensor>
checkpoint_layer_norm(const Tensor& input, const Tensor& weight, const Tensor& bias, long M, long N, double eps) {
rematerialize_function_t rt =
[=](const Tensors& vec) -> Tensors {
auto ret = at::native_layer_norm(vec.at(0), vec.at(1), vec.at(2), M, N, eps);
return {std::get<0>(ret), std::get<1>(ret), std::get<2>(ret)};
};
auto ret = CheckpointTensorImpl::make("native_layer_norm", rt, {input, weight, bias});
return {ret[0], ret[1], ret[2]};
}

std::tuple<Tensor, Tensor, Tensor>
checkpoint_layer_norm_backward(const Tensor& grad_out, const Tensor& input, const Tensor& mean, const Tensor& rstd, const Tensor& weight, long M, long N, std::array<bool, 3ul> output_mask) {
rematerialize_function_t rt =
[=](const Tensors& vec) -> Tensors {
auto ret = at::native_layer_norm_backward(vec.at(0), vec.at(1), vec.at(2), vec.at(3), vec.at(4), M, N, output_mask);
return {std::get<0>(ret), std::get<1>(ret), std::get<2>(ret)};
};
auto ret = CheckpointTensorImpl::make("native_layer_norm_backward", rt, {grad_out, input, mean, rstd, weight});
return {ret[0], ret[1], ret[2]};
}

bool checkpoint_equal(const Tensor& self, const Tensor& other) {
// there can't possibly be a reason to rematerialize
// a single bool so we'll just compute it now
return at::equal(decheckpoint(self), decheckpoint(other));
}

Scalar checkpoint__local_scalar_dense(at::Tensor const& a) {
return at::_local_scalar_dense(decheckpoint(a));
}
6 changes: 6 additions & 0 deletions aten/src/ATen/native/native_functions.yaml
Original file line number Diff line number Diff line change
@@ -688,6 +688,7 @@
dispatch:
CPU: bmm_cpu
CUDA: bmm_cuda
Checkpoint: checkpoint_bmm
supports_named_tensor: True

- func: bmm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!)
@@ -1657,11 +1658,13 @@
dispatch:
CPU: layer_norm_cpu
CUDA: layer_norm_cuda
Checkpoint: checkpoint_layer_norm

- func: native_layer_norm_backward(Tensor grad_out, Tensor input, Tensor mean, Tensor rstd, Tensor? weight, int M, int N, bool[3] output_mask) -> (Tensor, Tensor, Tensor)
dispatch:
CPU: layer_norm_backward_cpu
CUDA: layer_norm_backward_cuda
Checkpoint: checkpoint_layer_norm_backward

- func: linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor
python_module: nn
@@ -2687,12 +2690,14 @@
CPU: softmax_cpu
CUDA: softmax_cuda
MkldnnCPU: mkldnn_softmax
Checkpoint: checkpoint__softmax

- func: _softmax_backward_data(Tensor grad_output, Tensor output, int dim, Tensor self) -> Tensor
use_c10_dispatcher: full
dispatch:
CPU: softmax_backward_cpu
CUDA: softmax_backward_cuda
Checkpoint: checkpoint__softmax_backward_data

- func: split.Tensor(Tensor(a) self, int split_size, int dim=0) -> Tensor(a)[]
variants: function, method
@@ -5387,6 +5392,7 @@
CPU: legacy::cpu::_th_equal
CUDA: legacy::cuda::_th_equal
QuantizedCPU: quantized_equal
Checkpoint: checkpoint_equal
supports_named_tensor: True

- func: pow.Tensor_Tensor_out(Tensor self, Tensor exponent, *, Tensor(a!) out) -> Tensor(a!)
8 changes: 5 additions & 3 deletions torch/nn/functional.py
Original file line number Diff line number Diff line change
@@ -3776,9 +3776,10 @@ def multi_head_attention_forward(query, # type: Tensor

q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
if k is not None:
k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
# extremely bizarre, but got errors of dimension mistmatch here and below unless I changed view -> reshape
k = k.contiguous().reshape(-1, bsz * num_heads, head_dim).transpose(0, 1)
if v is not None:
v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
v = v.contiguous().reshape(-1, bsz * num_heads, head_dim).transpose(0, 1)

if static_k is not None:
assert static_k.size(0) == bsz * num_heads
@@ -3825,7 +3826,8 @@ def multi_head_attention_forward(query, # type: Tensor

attn_output = torch.bmm(attn_output_weights, v)
assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
# also had to change view to reshape
attn_output = attn_output.transpose(0, 1).contiguous().reshape(tgt_len, bsz, embed_dim)
attn_output = linear(attn_output, out_proj_weight, out_proj_bias)

if need_weights: