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.

Sign up for GitHub

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

Jump to bottom

[MPS] Fix batchnorm forward and backward pass #94351

Closed
wants to merge 2 commits into from
Closed

[MPS] Fix batchnorm forward and backward pass #94351

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4c94a24
Fix batchnorm forward and backward pass
DenisVieriu97 Feb 7, 2023
604b8ad
Fix layer norm
DenisVieriu97 Feb 9, 2023
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
64 changes: 49 additions & 15 deletions aten/src/ATen/native/mps/operations/Normalization.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,9 @@ void get_shapes(MPSShape* input_shape_readonly,
+ std::to_string(momentum) + ":" + std::to_string(train) + ":"
+ std::to_string(has_running_mean) + ":"
+ std::to_string(has_weight) + ":" + std::to_string(has_bias) + ":"
+ [ns_shape_key UTF8String] + ":" + native_mps::getMPSTypeString(self.scalar_type());
+ [ns_shape_key UTF8String] + ":"
+ native_mps::getTensorsStringKey({
self, weight_opt.value_or(Tensor()), bias_opt.value_or(Tensor()), running_mean_opt.value_or(Tensor()), running_var_opt.value_or(Tensor())});
auto input_mps_dtype = native_mps::getMPSDataType(self.scalar_type());
CachedGraph* cachedGraph = static_cast<CachedGraph *>(cache_->LookUp(key));

Expand Down Expand Up @@ -179,6 +181,7 @@ void get_shapes(MPSShape* input_shape_readonly,

MPSGraphTensor* updatedRunningMeanTensor = nil;
MPSGraphTensor* updatedRunningVarTensor = nil;
MPSGraphTensor *scaledInverseSqrtVariance = nil;

/*
If train:
Expand All @@ -194,6 +197,7 @@ Check if running mean exists (maybe do this check before making graph)

Compute the batch norm output and stats to be saved
*/
MPSGraphTensor *varTensor = nil;

if(train) {
// Compute mean and variance of the current batch
Expand All @@ -203,6 +207,7 @@ Check if running mean exists (maybe do this check before making graph)
MPSGraphTensor* batchVarianceTensor = [mpsGraph varianceOfTensor:inputTensor
axes:axes
name:nil];
varTensor = batchVarianceTensor;
if(has_running_mean) {
// TODO: This is not the formula used in PyTorch, is this OK? Seems more robust
// float besselCorrectionTerm = float(N) / std::max(N - 1.0f, 1.0f);
Expand Down Expand Up @@ -239,27 +244,41 @@ Check if running mean exists (maybe do this check before making graph)
updatedRunningVarTensor = [mpsGraph additionWithPrimaryTensor:scaledCorrectedBatchVar
secondaryTensor:scaledRunningVar
name:nil];
// Update saved mean and inverse std tensor
saveMeanTensor = batchMeanTensor;
saveVarTensor = batchVarianceTensor;
}
else {
saveMeanTensor = batchMeanTensor;
saveVarTensor = batchVarianceTensor;
}
// Update saved mean and inverse std tensor
MPSGraphTensor *epsilonTensor = [mpsGraph constantWithScalar:(double)epsilon
shape:@[@1]
dataType:MPSDataTypeFloat32];

MPSGraphTensor *varianceEps = [mpsGraph additionWithPrimaryTensor:batchVarianceTensor
secondaryTensor:epsilonTensor
name:@"varianceEps"];

MPSGraphTensor *sqrtVariance = [mpsGraph squareRootWithTensor:varianceEps
name:@"sqrtVariance"];
float primary = 1.0f;
MPSGraphTensor *primaryTensor = [mpsGraph constantWithScalar:primary dataType:MPSDataTypeFloat32];

scaledInverseSqrtVariance = [mpsGraph divisionWithPrimaryTensor:primaryTensor
secondaryTensor:sqrtVariance
name:nil];
Comment on lines +259 to +264
Copy link
Contributor

@malfet malfet Mar 7, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a reason not to use reciprocalWithTensor: here instead? It should be faster, shouldn't it? Or does it come with performance implications?
(PR is coming)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With the precise mode being default, there won't be any major perf implications.

// Update saved mean and inverse std tensor
saveMeanTensor = batchMeanTensor;
saveVarTensor = scaledInverseSqrtVariance;
}
else { // Test
TORCH_CHECK(has_running_mean);
saveMeanTensor = [mpsGraph identityWithTensor:runningMeanTensor
name:nil];
saveVarTensor = [mpsGraph identityWithTensor:runningVarTensor
name:nil];
varTensor = saveVarTensor;
}

// Compute output of batch norm
MPSGraphTensor* outputTensor = [mpsGraph normalizationWithTensor:inputTensor
meanTensor:saveMeanTensor
varianceTensor:saveVarTensor
varianceTensor:varTensor
gammaTensor:weightTensor
betaTensor:biasTensor
epsilon:(float)epsilon
Expand Down Expand Up @@ -351,6 +370,10 @@ Check if running mean exists (maybe do this check before making graph)

}

if(!train) {
save_mean.resize_({0});
save_var.resize_({0});
}
return std::tuple<Tensor&, Tensor&, Tensor&>(output, save_mean, save_var);
}

Expand Down Expand Up @@ -649,11 +672,24 @@ string get_mem_string(c10::MemoryFormat memory_format) {

if(train) {
// Use save_mean and save_var
float primary = 1.0f;
MPSGraphTensor *primaryTensor = [mpsGraph constantWithScalar:primary dataType:MPSDataTypeFloat32];
MPSGraphTensor *epsilonTensor = [mpsGraph constantWithScalar:(float)epsilon dataType:MPSDataTypeFloat32];
MPSGraphTensor *revertSaveVarTensor = saveVarTensor;
revertSaveVarTensor = [mpsGraph divisionWithPrimaryTensor: primaryTensor
secondaryTensor: revertSaveVarTensor
name: nil];
revertSaveVarTensor = [mpsGraph multiplicationWithPrimaryTensor: revertSaveVarTensor
secondaryTensor: revertSaveVarTensor
name: nil];
revertSaveVarTensor = [mpsGraph subtractionWithPrimaryTensor: revertSaveVarTensor
secondaryTensor: epsilonTensor
name: nil];
if(grad_input_mask[1]) {
gradWeightTensor = [mpsGraph normalizationGammaGradientWithIncomingGradientTensor:gradOutputTensor
sourceTensor:inputTensor
meanTensor:saveMeanTensor
varianceTensor:saveVarTensor
varianceTensor:revertSaveVarTensor
reductionAxes:axes
epsilon:(float)epsilon
name:nil];
Expand All @@ -668,7 +704,7 @@ string get_mem_string(c10::MemoryFormat memory_format) {
gradInputTensor = [mpsGraph normalizationGradientWithIncomingGradientTensor:gradOutputTensor
sourceTensor:inputTensor
meanTensor:saveMeanTensor
varianceTensor:saveVarTensor
varianceTensor:revertSaveVarTensor
gammaTensor:weightTensor
gammaGradientTensor:gradWeightTensor
betaGradientTensor:gradBiasTensor
Expand Down Expand Up @@ -890,8 +926,6 @@ string get_mem_string(c10::MemoryFormat memory_format) {
at::Tensor mean = std::get<1>(outputs);
at::Tensor variance = std::get<2>(outputs);

at::Tensor rstd = at::rsqrt(at::add(variance, eps));

std::vector<int64_t> stat_shape;
for (const auto idx : c10::irange(axis)) {
stat_shape.push_back(input_shape[idx]);
Expand All @@ -901,8 +935,8 @@ string get_mem_string(c10::MemoryFormat memory_format) {
stat_shape.push_back(1);
}
mean = mean.view(stat_shape);
rstd = rstd.view(stat_shape);
return std::make_tuple(out, mean, rstd);
variance = variance.view(stat_shape);
return std::make_tuple(out, mean, variance);
}

std::tuple<Tensor, Tensor, Tensor> layer_norm_backward_mps(
Expand Down
8 changes: 8 additions & 0 deletions test/test_mps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8785,6 +8785,8 @@ class TestConsistency(TestCase):
'nn.functional.bilinear': ['f32'],
'linalg.solve_triangular': ['f32'],
'triangular_solve': ['f32'],
'_native_batch_norm_legit': ['f32'],
'native_batch_norm': ['f32'],
}


Expand Down Expand Up @@ -8968,6 +8970,9 @@ class TestConsistency(TestCase):
'zero_': ['f16', 'f32'],
'linalg.solve_triangular': ['f32'],
'triangular_solve': ['f32'],
'_native_batch_norm_legit': ['f32'],
'native_batch_norm': ['f32'],
'native_layer_norm': ['f32'],
}

# These ops that are problematic. So never run them even when
Expand Down Expand Up @@ -9183,6 +9188,9 @@ def get_samples():
elif (op.name == "masked.mean"):
atol = 7e-4
rtol = 2e-3
elif (op.name == "native_layer_norm"):
atol = 1e-4
rtol = 1.3e-5
else:
atol = None
rtol = None
Expand Down