KD example fix for new torch/hf causing FDSP save error (NVIDIA#645)

## What does this PR do?

**Type of change:** ? Bug Fix

**Overview:** ? `llm_distill` example was hanging during save since
somehow now the weights on other ranks are being deleted during
`model.export()` too early. Fixed via synchronizing the processes
beforehand.

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->

Signed-off-by: Asha Anoosheh <[email protected]>

Author: AAnoosheh

examples/llm_distill/requirements.txt

@@ -1,3 +1,4 @@
 pyarrow
+torchao>=0.14.1
 transformers<5.0
 trl>=0.23.0

modelopt/torch/distill/plugins/huggingface.py

@@ -64,13 +64,13 @@ def save_model(
             output_dir = self.args.output_dir
         model = self.accelerator.unwrap_model(self.model)
         if not _internal_call and self.is_fsdp_enabled:
-            state_dict = self.accelerator.get_state_dict(self.model)
+            with model.hide_teacher_model(enable=export_student):
+                state_dict = self.accelerator.get_state_dict(self.model)
             modelopt_state = mto.modelopt_state(model)
             if export_student:
+                # Need to wait, otherwise FSDP weights may be deleted before rank 0 can gather them
+                self.accelerator.wait_for_everyone()
                 model = model.export()
-                # remove teacher model from state dict since FSDP forces
-                # expose_minimal_state_dict to be False
-                state_dict = {k: v for k, v in state_dict.items() if "_teacher_model" not in k}
 
             if self.accelerator.is_main_process:
                 model.save_pretrained(