Add pre-training transform to convert BatchNorm to BatchNormInternal

[pranav-prakash]

Description: As opset 14 lacks saved_mean/saved_inv_std, we have to convert nodes to a custom op if we want to ensure good performance on backward pass. This is the second in a series of PRs to implement BN training on CPU (first was #6946).

As part of this we also remove the existing AdjustBatchNormOutputs pass – this only works for opset 9 BN and is subsumed by the more general BatchNormReplacement we define in this PR.

WIP – I think the only thing missing is a unit test. I'd also appreciate any feedback on whether it's necessary to copy/paste the entire schema-def for the BatchNormInternal op or whether something like protobuf's toBuilder exists for these schema defs.

[SherlockNoMad]

Hi @pranav-prakash, thanks for contributing this !

Let us know when this is ready for review. I will help you merging this change.

[pranav-prakash]

@SherlockNoMad
Thanks, I just added the test so it should be ready for review. One question I had was that the test required me to explicitly import the com.microsoft domain in order for it to recognize the BatchNormInternal op – I'm not sure if this is expected behavior. When loading a model from disk is com.microsoft automatically added as a valid import?

[SherlockNoMad]

I think we also need to update gradient_builder_registry.cc, to register the gradient builder for BatchNormInternal.
We should be able to reuse "GetBatchNormalizationGradient" though.

[SherlockNoMad]

@SherlockNoMad
Thanks, I just added the test so it should be ready for review. One question I had was that the test required me to explicitly import the com.microsoft domain in order for it to recognize the BatchNormInternal op – I'm not sure if this is expected behavior. When loading a model from disk is com.microsoft automatically added as a valid import?

I can imagine msDomain is needed to make the UT work. In the training process, msDomain kernel registry will be automatically included at program start time, coz most of the gradient ops are under MSDomain.

[SherlockNoMad]

Also need to add BatchNormInternal into the STOP_GRADIENT_EDGES in gradient_graph_builder.h

[SherlockNoMad]

Thanks again for the PR. The bulk part is looking good!
We should be good to merge once the comments are addressed

[pranav-prakash]

@SherlockNoMad
Thanks for the quick review! I addressed the comments you left.

register gradient builder for BatchNormInternal.

Done. I'll add a CPU implementation of the gradient in a subsequent PR.

I'll also register an implementation for BatchNormInternal (re-using existing BatchNorm kernel) in a subsequent PR. I believe there's also a pre-existing CUDA kernel for batch norm training & gradient, but it seems to be broken (tests were disabled) so someone else who's both more familiar and capable of testing it will have to modify those parts.

[SherlockNoMad]

/azp run Linux CPU CI Pipeline,Linux CPU x64 NoContribops CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,MacOS CI Pipeline,MacOS NoContribops CI Pipeline,Windows CPU CI Pipeline,Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline

[SherlockNoMad]

/azp run orttraining-linux-ci-pipeline,orttraining-mac-ci-pipeline,orttraining-linux-gpu-ci-pipeline,centos7_cpu,Linux CPU Minimal Build E2E CI Pipeline,Linux Nuphar CI Pipeline,MacOS NoContribops CI Pipeline,Linux OpenVINO CI Pipeline,orttraining-distributed

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 8 pipeline(s).

[SherlockNoMad]

sure. Appreciate it if you can contribute the CPU BatchNormGradient kernel !!

[pranav-prakash]

@SherlockNoMad
Seems like maybe we can't remove InsertMaxPoolOutput in gradient_graph_builder.cc after all? It fails the test

1: C++ exception with description "/onnxruntime_src/orttraining/orttraining/core/graph/gradient_builder_base.h:102 onnxruntime::training::ArgDef onnxruntime::training::GradientBuilderBase::O(size_t) const i < node_->OutputDefs().size() was false. 

although I'm not sure why this is since we register it in graph_transformer_utils.cc as well. I'm not sure about the order in which these things are called, but is it possible that the tests only call into gradient_graph_builder.cc but not graph_transformer_utils.cc?

[SherlockNoMad]

ah... ic.. the GradientCheckerTest.MaxPoolGrad is not invoking the GeneratePreTrainingTransformers, as it's taking the InferenceSession path to construct the test model. and InferenceSession's default graph transformer doesn't include InsertMaxpoolOutput...

Please revert the Maxpool related changes, we will fix this in the latter PRs. Sorry for the inconvenience caused.

[pranav-prakash]

@SherlockNoMad Done, thanks for determining the cause!

[pranav-prakash]

the GradientCheckerTest.MaxPoolGrad is not invoking the GeneratePreTrainingTransformers

Actually this affects the unit tests for BatchNormGradient as well (seen in the below draft PR). Even if we register

rule_based_graph_transformer->Register(std::make_unique<BatchNormReplacement>());

for the constructor of GradientGraphBuilder, the reachable_nodes_ and forward_reachable_nodes are determined before the transform is run. So the unit test for BatchNormGradient will fail with the message

Cannot compute the partial derivative for 'input1' as it's unreachable from the output node(s).

since reachable_nodes_ has the node before transformation (BatchNormalization) but the CheckNodeArgsReachable is done after the graph transformation. I've worked around it for now by moving the graph transform before the transitive closure computation, but as you mentioned the cleaner solution is to just invoke GeneratePreTrainingTransformers directly.

[SherlockNoMad]

/azp run Linux CPU CI Pipeline,Linux CPU x64 NoContribops CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,MacOS CI Pipeline,MacOS NoContribops CI Pipeline,Windows CPU CI Pipeline,Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline

[SherlockNoMad]

/azp run orttraining-linux-ci-pipeline,orttraining-mac-ci-pipeline,orttraining-linux-gpu-ci-pipeline,centos7_cpu,Linux CPU Minimal Build E2E CI Pipeline,Linux Nuphar CI Pipeline,MacOS NoContribops CI Pipeline,Linux OpenVINO CI Pipeline,orttraining-distributed

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 8 pipeline(s).

[SherlockNoMad]

/azp run Windows WebAssembly CI Pipeline, orttraining-amd-gpu-ci-pipeline, orttraining-ortmodule, orttraining-ortmodule-distributed

[azure-pipelines]

Azure Pipelines successfully started running 4 pipeline(s).