I might have made a mistake when copying the code from the script to post here  In the code that produces the behavior mentioned in the OP, the model is wrapped actually exactly like you propsed, with the addition of the find_unused_parameters argument:
model = torch.nn.parallel.DistributedDataParallel(model, 
                                                  device_ids=[args.local_rank],
                                                  output_device=args.local_rank,
                                                  find_unused_parameters=True)

(Edited the previous post accordingly.)