text: Title: Error restarting training and inability to enable checkpoint activations after compiling Fairseq model using torch.compile()
  Body: ## ❓ Questions and Help
  
  
  #### compiled fairseq model using torch.compile() and observed significant training speed . but observed few issues . 
  1) Error occurred while restarting the training from checkpoint_last .
  2) not able to enable  --checkpoint-activations while using torch.compile 
  
  #### Code
  
    ```
  if cfg.distributed_training.ddp_backend == "fully_sharded":
  
          with fsdp_enable_wrap(cfg.distributed_training):
  
              model = fsdp_wrap(task.build_model(cfg.model))
  else:
  
          model = task.build_model(cfg.model)
  
          model = torch.compile(model)
  ```
   
  
  #### compiled model using torch.compile() ,enabled checkpoint activations using --checkpoint-activations and observed below error.
  KeyError: packed_non_tensor_outputs                                                                                                                                                                        
                                                                                                                                                                                                             
  from user code:                                                                                                                                                                                            
     File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/modules/checkpoint_activations.py", line 67, in <graph break in _checkpointed_forward>                                
      packed_non_tensor_outputs = parent_ctx_dict["packed_non_tensor_outputs"]     
  
  Also when I tried to restart training from the last checkpoint, the below error occurred.
  
  2023-05-02 18:08:17 | INFO | fairseq.trainer | Preparing to load checkpoint checkpoint/checkpoint_last.pt
  Traceback (most recent call last):
    File "/home/santha-11585/miniconda3/envs/ddp/bin/fairseq-train", line 8, in <module>
      sys.exit(cli_main())
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq_cli/train.py", line 632, in cli_main
      distributed_utils.call_main(cfg, main)
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/distributed/utils.py", line 344, in call_main
      torch.multiprocessing.spawn(
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 239, in spawn
      return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 197, in start_processes
      while not context.join():
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 160, in join
      raise ProcessRaisedException(msg, error_index, failed_process.pid)
  torch.multiprocessing.spawn.ProcessRaisedException: 
  
  -- Process 6 terminated with the following error:
  Traceback (most recent call last):
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/trainer.py", line 567, in load_checkpoint
      self.model.load_state_dict(
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/distributed/module_proxy_wrapper.py", line 53, in load_state_dict
      return self.module.module.load_state_dict(*args, **kwargs)
  TypeError: Module.load_state_dict() got an unexpected keyword argument 'model_cfg'
  
  During handling of the above exception, another exception occurred:
  
  Traceback (most recent call last):
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
      fn(i, *args)
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/distributed/utils.py", line 328, in distributed_main
      main(cfg, **kwargs)
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq_cli/train.py", line 170, in main
      extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/checkpoint_utils.py", line 248, in load_checkpoint
      extra_state = trainer.load_checkpoint(
    File "/home/santha-11585/miniconda3/envs/ddp/lib/python3.10/site-packages/fairseq/trainer.py", line 579, in load_checkpoint
      raise Exception(
  Exception: Cannot load model parameters from checkpoint checkpoint/checkpoint_last.pt; please ensure that the architectures match.
  
                                                                                                                            
                                                                                                                                                                                                             
  
  
  #### What's your environment?
  
   - fairseq Version (e.g., 1.0 or main): 0.12.2
   - PyTorch Version (e.g., 1.0) : 2.0.0+cu117
   - OS (e.g., Linux): ubuntu
   - How you installed fairseq (`pip`, source): yes
   - Build command you used (if compiling from source):
   - Python version:3.10
   - CUDA/cuDNN version: 11.7
   - GPU models and configuration: A6000
   - Any other relevant information:
  