fix: invalidate ZeRO-3 param coordinator trace in add_hooks

[roycho96]

What does this PR do?

This PR fixes a IndexError: pop from an empty deque error that occurs when using GKDTrainer with DeepSpeed ZeRO-3.

Problem

When using ZeRO-3, training fails on the second step with:

File "deepspeed/runtime/zero/partitioned_param_coordinator.py", line 217, in record_parameters
    step_id = self.__step_id_module_fetched_for[sub_module.ds_id].popleft()
IndexError: pop from an empty deque

Root Cause

DeepSpeed registers fwd_pre_hook separately from forward_hooks list. The current remove_hooks() doesn't remove fwd_pre_hook, so during generation, reset_step() is still called on every forward pass, corrupting the PartitionedParameterCoordinator's internal trace state.

When add_hooks() restores hooks without resetting this corrupted state, the subsequent training forward fails.

Solution

Call coordinator._invalidate_trace() in add_hooks() before re-registering hooks to reset the coordinator to a clean state.

Testing

• Tested manually with GKDTrainer + DeepSpeed ZeRO-3
• Config: lmbda=0.7, seq_kd=True
• Model: Llama 3.2 1B (student model), Llama 3.2 3B (teacher model)
• DeepSpeed version: 0.17.6

lmbda	seq_kd	Before fix	After fix
0.0	True	✅ Pass	✅ Pass
0.7	True	❌ Fail	✅ Pass

This fix also benefits other trainers using unwrap_model_for_generation with ZeRO-3

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[kashif]

thanks @roycho96 checking!

[qgallouedec]

each time you use a getattr/hasattr, can you please explicitly and clearly mention via a comment in which case the object doesn't have the attribute? Eg:

if not hasattr(model, "optimizer"):  # before the first training step, the model has no optimizer

[roycho96]

Thanks for the review! Updated the code.
I removed unnecessary hasattr/getattr checks for is_invalid_trace() and _invalidate_trace() since these methods always exist in PartitionedParameterCoordinator.
Also, I added a comment that param_coordinator only exists in ZeRO stage 3.