def _collect_main_grad_data_for_unscaling(self):
        """
        Note: this should be equivalent to the float-16 optimizer's method,
        but writtent differently, so the two should be combined.
        """
        return [
            param.grad.data
            for group in self.optimizer.param_groups
            for param in group["params"]
        ]

        Note: this should be equivalent to the float-16 optimizer's method,
        but writtent differently, so the two should be combined.
        
training/DeepSpeed-Domino/domino/optimizer/distrib_optimizer.py