enable weight cleanup (#289)

Author: allenwang28

.gitignore

@@ -197,6 +197,5 @@ wandb/
 assets/wheels/vllm*.whl
 
 # DCP artifacts
-model_state_dict/
 forge_dcp_tmp/
 demo_top_down.md

apps/grpo/main.py

@@ -455,9 +455,9 @@ async def continuous_training():
                 await policy.update_weights.fanout(training_step)
                 t.step("update_weights")
 
-                # if training_step >= 2:
-                #     await drop_weights(training_step - 1)
-                #     t.step("drop_weights")
+                if training_step >= 2:
+                    await drop_weights(training_step - 1)
+                    t.step("drop_weights")
 
                 t.stop()
                 restart_tracer = True