Merge: [DLRM/PyT] 21.10 container update with added graph support and BYOD capability

Author: nv-kkudrynski

PyTorch/Recommendation/DLRM/Dockerfile

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:21.04-py3
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:21.10-py3
 FROM ${FROM_IMAGE_NAME}
 
 ADD requirements.txt .

PyTorch/Recommendation/DLRM/README.md

@@ -196,7 +196,6 @@ esac
 ################################################################################
 
 if [ "${#numactl_args[@]}" -gt 0 ] ; then
-    set -x
     exec numactl "${numactl_args[@]}" -- "${@}"
 else
     exec "${@}"

PyTorch/Recommendation/DLRM/bind.sh

@@ -14,28 +14,26 @@
 
 import torch
 from torch.autograd import Function
-from apex import amp
+
 
 if torch.cuda.get_device_capability()[0] >= 8:
-    print('Using the Ampere-optimized dot interaction kernels')
     from dlrm.cuda_ext import interaction_ampere as interaction
 else:
-    print('Using the Volta-optimized dot interaction kernels')
     from dlrm.cuda_ext import interaction_volta as interaction
 
 
 class DotBasedInteract(Function):
     """ Forward and Backward paths of cuda extension for dot-based feature interact."""
 
     @staticmethod
-    @amp.half_function
+    @torch.cuda.amp.custom_fwd(cast_inputs=torch.half)
     def forward(ctx, input, bottom_mlp_output):
         output = interaction.dotBasedInteractFwd(input, bottom_mlp_output)
         ctx.save_for_backward(input)
         return output
 
     @staticmethod
-    @amp.half_function
+    @torch.cuda.amp.custom_bwd
     def backward(ctx, grad_output):
         input, = ctx.saved_tensors
         grad, mlp_grad = interaction.dotBasedInteractBwd(input, grad_output)