[Enhance] Add layers[-1] to norm & lm_head prefetch

Author: nil0x9

xtuner/v1/model/dense/dense.py

@@ -270,6 +270,7 @@ def fully_shard(
             offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else OffloadPolicy(),
         )
         self.set_modules_to_forward_prefetch([self.embed_tokens, self.layers["0"]])  # type: ignore
+        list(self.layers.values())[-1].set_modules_to_forward_prefetch([self.norm, self.lm_head])  # type: ignore
 
         self._to_empty_meta()
 

xtuner/v1/model/moe/moe.py

@@ -780,6 +780,7 @@ def fully_shard(
             offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else OffloadPolicy(),
         )
         self.set_modules_to_forward_prefetch([self.embed_tokens, self.layers["0"]])  # type: ignore
+        list(self.layers.values())[-1].set_modules_to_forward_prefetch([self.norm, self.lm_head])  # type: ignore
 
         for _, module in self.named_modules():
             if isinstance(module, nn.Embedding):