Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -115,6 +115,12 @@ class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
                 new_block.norm_topk_prob = orig_block.norm_topk_prob
                 layer.mlp = new_block
                 print(type(layer.mlp))
+        # 释放预训练模型内存
+        del pretrained_model
+        import gc
+        gc.collect()
+        torch.cuda.empty_cache()
+        print("原始预训练模型已释放")
 
 def main():
     config = DenseBackwardOLMoEConfig(        # 官方模型参数