Repository: hpcaitech/ColossalAI
File: colossalai/zero/gemini/gemini_ddp.py (Python)

Code before:
            error_params = ["Reduction failed at followed parameters:"]
            for param in self.param2name:
                if not is_ddp_ignored(param) and not getattr(param, "_gemini_reduced"):
                    error_params.append(self.param2name[param])
            error_str = "\n\t".join(error_params)
            raise RuntimeError(
                "ZERO DDP error: the synchronization of gradients doesn't exit properly.",
                "The most possible reason is that the model is not compatible with GeminiDDP.\n",
                f"{error_str}",
            )
        self._setup_grads_ptr()
        if self.enable_gradient_accumulation and not self.accumulating_grads:
            self.accumulating_grads = True  # Turn on the state of gradient accumulation.
        self._logger.debug(
            f"comp cuda demand time: {self.gemini_manager._comp_cuda_demand_time}, layout time: {self.gemini_manager._layout_time}, evict time: {self.gemini_manager._evict_time}, CPU->CUDA vol: {self.gemini_manager._h2d_volume}B, CUDA->CPU vol: {self.gemini_manager._d2h_volume}"
        )
        self.gemini_manager.post_iter()

    def backward(self, loss: torch.Tensor):
        self._pre_backward()
        with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook):
            loss.backward()
        self._post_backward()

    def backward_by_grad(self, tensor, grad):
        self._pre_backward()
        with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook):
            torch.autograd.backward(tensor, grad)
        self._post_backward()

    def grad_handle(self, p, grad):
        setattr(p, "_gemini_reduced", True)
        empty_grad = torch.empty_like(grad)
        free_storage(empty_grad)
        with torch._C.DisableTorchFunction():
            chunk = self.chunk_manager.get_chunk(p)
            if chunk.tensors_info[p].state != TensorState.HOLD_AFTER_BWD:
                raise RuntimeError(
                    f"Parameter `{self.param2name[p]}` failed at the gradient reduction. "
                    "Some unsupported torch function is operated upon this parameter."
                )
            grad_chunk = chunk
            if not self.reuse_fp16_chunk:
                if not self.accumulating_grads:
                    grad_chunk = self.chunk_manager.init_grad_chunk(chunk)
                else:
                    assert chunk.grad_chunk is not None
                    if chunk.grad_chunk not in self.chunk_manager.accessed_chunks:
                        grad_chunk = self.chunk_manager.rearrange_accumulated_grad_chunk(chunk)
                    else:
                        grad_chunk = chunk.grad_chunk

Review comment (bug):
`backward_by_grad()` shoud not be called, since Gemini is not compatible with pipeline. Raise a runtime error here.

Code after:
            error_params = ["Reduction failed at followed parameters:"]
            for param in self.param2name:
                if not is_ddp_ignored(param) and not getattr(param, "_gemini_reduced"):
                    error_params.append(self.param2name[param])
            error_str = "\n\t".join(error_params)
            raise RuntimeError(
                "ZERO DDP error: the synchronization of gradients doesn't exit properly.",
                "The most possible reason is that the model is not compatible with GeminiDDP.\n",
                f"{error_str}",
            )
        self._setup_grads_ptr()
        if self.enable_gradient_accumulation and not self.accumulating_grads:
            self.accumulating_grads = True  # Turn on the state of gradient accumulation.
        self._logger.debug(
            f"comp cuda demand time: {self.gemini_manager._comp_cuda_demand_time}, layout time: {self.gemini_manager._layout_time}, evict time: {self.gemini_manager._evict_time}, CPU->CUDA vol: {self.gemini_manager._h2d_volume}B, CUDA->CPU vol: {self.gemini_manager._d2h_volume}"
        )
        self.gemini_manager.post_iter()

    def backward(self, loss: torch.Tensor):
        self._pre_backward()
        with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(self.param_op_hook):
            loss.backward()
        self._post_backward()

    def backward_by_grad(self, tensor, grad):
        raise RuntimeError("Gemini is not compatible with pipeline. backward_by_grad shoudn't be called in Gemini.")

    def grad_handle(self, p, grad):
        setattr(p, "_gemini_reduced", True)
        empty_grad = torch.empty_like(grad)
        free_storage(empty_grad)
        with torch._C.DisableTorchFunction():
            chunk = self.chunk_manager.get_chunk(p)
            if chunk.tensors_info[p].state != TensorState.HOLD_AFTER_BWD:
                raise RuntimeError(
                    f"Parameter `{self.param2name[p]}` failed at the gradient reduction. "
                    "Some unsupported torch function is operated upon this parameter."
                )
            grad_chunk = chunk
            if not self.reuse_fp16_chunk:
                if not self.accumulating_grads:
                    grad_chunk = self.chunk_manager.init_grad_chunk(chunk)
                else:
                    assert chunk.grad_chunk is not None
                    if chunk.grad_chunk not in self.chunk_manager.accessed_chunks:
                        grad_chunk = self.chunk_manager.rearrange_accumulated_grad_chunk(chunk)
                    else:
                        grad_chunk = chunk.grad_chunk
                        chunk.grad_chunk.l2_norm = None

                # hold -> compute -> hold after bwd