记一次刻苦铭心的deepspeed调试经历

背景

笔者最近加入了学校的ai课题组，大一的暑假比较空闲，在做一个师兄丢给我的3D diffusion相关的项目。其中设计需要对 flux kontext加模块微调。 但是flux kontext光是推理就需要用掉40G显存，而课题组的计算服务器只有48GB的A6000， 显存资源及其有限。如果直接在单卡开train会直接CUDA OUT OF MEMORY。

经过

师兄给笔者介绍了deepspeed 的activation checkpoint技术，可以节省显存，笔者一开始没有经验，按着师兄的代码哐哐乱抄，把整个可训练模块放在了checkpoint里面，同时模仿师兄的代码在函数返回值的最后添加require_grad_(True),代码如下：

def forward(self,hidden_states: torch.Tensor,encoder_hidden_states: torch.Tensor,temb: torch.Tensor,image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,joint_attention_kwargs: Optional[Dict[str, Any]] = None):encoder_hidden_states, hidden_states = self._forward(hidden_states,encoder_hidden_states,temb,image_rotary_emb,joint_attention_kwargs)return encoder_hidden_states.require_grad_(True), hidden_states.require_grad_(True)

def _forward( self, hidden_states: torch.Tensor, # ... the rest of the file is unchanged ...

 

触发报错：

File "/usr/local/lib/python3.10/dist-packages/torch/functional.py", line 1626, in norm
return torch.linalg.vector_norm(input, _p, _dim, keepdim, dtype=dtype)
RuntimeError: linalg.vector_norm: Expected a floating point or complex tensor as input. Got Long。

感到很莫名其妙，一开始以为定义参数出了问题，翻遍整个项目库也没有找到torch.long的可训练参数

后来跟着ai 排查了n轮后问ai说是不能这么搞

然后去掉了require_grad_

 

def forward(self,hidden_states: torch.Tensor,encoder_hidden_states: torch.Tensor,temb: torch.Tensor,image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,joint_attention_kwargs: Optional[Dict[str, Any]] = None):encoder_hidden_states, hidden_states = self._forward(hidden_states,encoder_hidden_states,temb,image_rotary_emb,joint_attention_kwargs)return encoder_hidden_states, hidden_states
def _forward( self, hidden_states: torch.Tensor, # ... the rest of the file is unchanged ...

 

 

结果直接报错说没有梯度可以进行反向传播， 

File "/usr/local/lib/python3.10/dist-packages/torch/autograd/graph.py", line 681, in _engine_run_backward   
return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

ai说是由于整个模型被包裹导致计算图丢失了

综合分析后就是不能把整个模型都放在checkpoint里面，如果整个放进去无论如何都会触发错误。

然后笔者就病急乱投医，把不可训练部分扔进checkpoint里面，但是根本没有提升效果。原因也很明显，因为不可训练部分本来就不会记录激活值。自然activation checkpoint也没有用了。

解决

今天笔者联想到flux里面总共有10几个Dit block，想到了新方法：把除了第一个block以外的block全部丢进checkpoint里面，这样既有了梯度可供backward ，又不会cuda out of memory，速度也还可以接受，赢！。

从开始动工到最终能跑，过了整整三周，但总算是把这个搞定了。

经验&吐槽

调试问题如果要问ai,必须把整个项目设计的代码提供给ai，这样ai才能给出针对性的方向。

同时也不能全信ai,现在ai的幻觉仍然很严重，必须亲自去查官方的文档（比如笔者这几天在deepspeed的官方文档翻了半天）

理论知识要提前具备，笔者写训练代码时没有搞懂flow matching,视图直接套用diffusion的训练框架训flow matching ，结局可想而知。后来去b站学了mit的公开课才算弄懂。

ai训练优化这一块其实也充斥者大量oi里面的时间换空间等类似的优化思想，从这一点来说，oi还是有点用的（

修改hugging face的模型的架构是真的屎，代码库超级庞大，各种类继承满天飞，要调试简直是噩梦。