Update FSDP, AC, Compile ordering

src/prime_rl/trainer/config.py

@@ -24,6 +24,15 @@ class ActivationCheckpointConfig(BaseModel):
     ] = 1
 
 
+class CompileConfig(BaseModel):
+    """Configures model compilation."""
+
+    fullgraph: Annotated[
+        bool,
+        Field(description="Whether to compile the transformer blocks with fullgraph."),
+    ] = False
+
+
 class ModelConfig(BaseConfig):
     """Configures the model for training."""
 
@@ -37,11 +46,11 @@ class ModelConfig(BaseConfig):
     attn: Annotated[AttnImplementation, Field(description="The attention implementation to use.")] = "flash_attention_2"
 
     compile: Annotated[
-        bool,
+        CompileConfig | None,
         Field(
             description="Whether to compile the model using `torch.compile`. Currently discouraged because it was found to destabilize training.",
         ),
-    ] = False
+    ] = None
 
     ac: Annotated[
         ActivationCheckpointConfig | None,

src/prime_rl/trainer/model.py

@@ -12,7 +12,7 @@
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 from transformers.tokenization_utils import PreTrainedTokenizer
 
-from prime_rl.trainer.config import ActivationCheckpointConfig, ModelConfig
+from prime_rl.trainer.config import ActivationCheckpointConfig, CompileConfig, ModelConfig
 from prime_rl.trainer.parallel_dims import ParallelDims
 from prime_rl.utils.logger import get_logger
 
@@ -76,17 +76,27 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
     mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
     # TODO: Support dp_replicate
     hsdp_mesh = parallel_dims.world_mesh["dp_shard_cp"]
-    for layer_id, transformer_block in enumerate(model.model.layers):
-        if config.reshard_after_forward:
-            layer_reshard_after_forward = layer_id < len(model.model.layers) - 1
-        else:
-            layer_reshard_after_forward = False
+
+    fully_shard(
+        model.model.embed_tokens,
+        mesh=hsdp_mesh,
+        mp_policy=mp_policy,
+        reshard_after_forward=config.reshard_after_forward,
+    )
+
+    for transformer_block in model.model.layers:
         fully_shard(
             transformer_block,
             mesh=hsdp_mesh,
             mp_policy=mp_policy,
-            reshard_after_forward=layer_reshard_after_forward,
+            reshard_after_forward=config.reshard_after_forward,
         )
+    fully_shard(
+        [model.lm_head, model.model.norm],
+        mesh=hsdp_mesh,
+        mp_policy=mp_policy,
+        reshard_after_forward=config.reshard_after_forward,
+    )
 
     fully_shard(model, mesh=hsdp_mesh, mp_policy=mp_policy, reshard_after_forward=config.reshard_after_forward)
 
@@ -115,19 +125,33 @@ def apply_ac(model: nn.Module, ac_config: ActivationCheckpointConfig):
         if layer_id % ac_config.freq == 0:
             transformer_block = checkpoint_wrapper(transformer_block, preserve_rng_state=False)
         model.model.layers.register_module(layer_name, transformer_block)
+    get_logger().info(f"Applied activation checkpointing (freq={ac_config.freq})")
+
+
+def apply_compile(model: nn.Module, compile_config: CompileConfig):
+    for layer_id, transformer_block in enumerate(model.model.layers):
+        model.model.layers[layer_id] = torch.compile(transformer_block, fullgraph=compile_config.fullgraph)
+    get_logger().info(f"Compiled {len(model.model.layers)} layers (fullgraph={compile_config.fullgraph})")
 
 
 def setup_model(config: ModelConfig, parallel_dims: ParallelDims) -> nn.Module:
-    model = get_model(config, device=torch.device("cpu"))
+    device = torch.device("cpu")
+    model = get_model(config, device=device)
+
+    # the right order is AC -> Compile -> FSDP
+    if config.ac is not None:
+        apply_ac(model, config.ac)
+    if config.compile is not None:
+        apply_compile(model, config.compile)
+
     setup_fsdp(model, config, parallel_dims)
+
+    # if device == torch.device("meta"):
     # TODO: This is used if the model is loaded with meta device to save cpu memory
     # However, the loading seems to be wrong as the loss and reward curves are different
     # load_dcp_from_hf(model, config)
-    if config.ac is not None:
-        apply_ac(model, config.ac)
-    if config.compile:
-        model = torch.compile(model)
-    # TODO: This should be type-hinted as FSDP version of the model
+    #     load_dcp_from_hf(model, config)
+
     return model