Remove moe_mlpwi for now

Author: abhinavgoel95

src/maxtext/configs/base.yml

@@ -317,7 +317,6 @@ mlpwi: 'remat'
 mlpwi_0: 'remat'
 mlpwi_1: 'remat'
 mlpwo: 'remat'
-moe_mlpwi: 'remat'
 moe_mlpwi_0: 'remat'
 moe_mlpwi_1: 'remat'
 moe_mlpwo: 'remat'

src/maxtext/configs/pyconfig_deprecated.py

@@ -516,7 +516,6 @@ def validate_and_assign_remat_tensors(keys):
       "mlpwi_0",
       "mlpwi_1",
       "mlpwo",
-      "moe_mlpwi",
       "moe_mlpwi_0",
       "moe_mlpwi_1",
       "moe_mlpwo",

src/maxtext/configs/types.py

@@ -896,10 +896,6 @@ class RematAndOffload(BaseModel):
       RematLocation.REMAT,
       description="Remat policy for the second MLP layer's output.",
   )
-  moe_mlpwi: RematLocation = Field(
-      RematLocation.REMAT,
-      description="Remat policy for the first MoE layer's intermediate output.",
-  )
   moe_mlpwi_0: RematLocation = Field(
       RematLocation.REMAT,
       description="Remat policy for the first part of a gated MoE's output.",