I don’t think there is a decent way to get rid of it in v1.6. One hacky solution might be introducing a tiny unused parameter in the model (e.g., self.unused = nn.Linear(1, 1)), and then set find_unused_parameters=True in DDP ctor, which would disable bucket rebuilt.

  
      github.com
  
  
    pytorch/pytorch/blob/v1.6.0/torch/csrc/distributed/c10d/reducer.cpp#L381 1


// Rebuild bucket only if 1) it is the first time to rebuild bucket 2)
// unused_parameters_ is empty, currently it does not support when there are
// unused parameters 3) this backward pass needs to run allreduce. Here, we
// just dump tensors and their parameter indices into rebuilt_params_ and
// rebuilt_param_indices_ based on gradient arriving order, and then at the
// end of finalize_backward(), buckets will be rebuilt based on
// rebuilt_params_ and rebuilt_param_indices_, and then will be broadcasted
// and intialized. Also we only need to dump tensors and parameter indcies of
// one replica.
if (!has_rebuilt_bucket_ && unused_parameters_.empty() &&
    index.replica_index == 0) {
  rebuilt_params_.push_back(
      replicas_[index.replica_index][index.variable_index]);
  rebuilt_param_indices_.push_back(index.variable_index);
}

// If there are model parameters that went unused when computing the model
// output, they won't be part of the autograd graph, and won't receive
// gradients. These parameters are discovered in the `prepare_for_backward`
// function and their indexes stored in the `unused_parameters_` vector.



  
  
    
    
  
  


This looks like a regression to me. @Yanli_Zhao has a some recent work to reduce DDP memory footprint and hopefully that can help.
BTW, if possible, can you try the same script with PyTorch v1.5, it will help to confirm if bucket reconstruction is indeed the culprit.