chore: revert accidental changes to attention.py and layers.py during rebase

Author: judyzhaoxinwu

axlearn/common/attention.py

@@ -2993,7 +2993,6 @@ def attention_thunk(target: Tensor) -> tuple[Optional[NestedTensor], Tensor]:
                 assert cached_states is not None
                 assert segment_ids is None
                 assert target_positions is None
-
                 atten_state, atten_output = self.attention.init_states(
                     time_step=cached_states["attention"],
                     query=target,
@@ -3004,7 +3003,6 @@ def attention_thunk(target: Tensor) -> tuple[Optional[NestedTensor], Tensor]:
                 assert cached_states is not None
                 assert segment_ids is None
                 assert target_positions is None
-
                 atten_state, atten_output = self.attention.extend_step(
                     cached_states["attention"],
                     target,

axlearn/common/layers.py

@@ -882,9 +882,7 @@ class Config(BaseLayer.Config):
     @classmethod
     def default_config(cls):
         cfg = super().default_config()
-        # Prevent OOM on 405b
-        # cfg.param_partition_spec = (None, "model")
-        cfg.param_partition_spec = ("fsdp", "model")
+        cfg.param_partition_spec = (None, "model")
         # By default, initialize to Gaussian with std=1/sqrt(dim), e.g., 0.036 when dim=768.
         #
         # This is the same as: