diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan
index 82084e7c..cfa631dd 160000
--- a/3rdparty/torchtitan
+++ b/3rdparty/torchtitan
@@ -1 +1 @@
-Subproject commit 82084e7c67c895e1f91b6d90887d3f7c83909577
+Subproject commit cfa631dd08919107867a65694da32b3faf4fd46d
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4627965c..46a45e3e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # Changelog
 
+## v0.1.0 [dev]
+
+- Changed
+  - Mixed precision is now handled by FSDP even if world_size=1.
+  - Dropped Instella-3B model config.
+
+
 ## v0.0.2 [dev]
 
 - Changed
diff --git a/alto/components/converter.py b/alto/components/converter.py
index 88ecb983..cd43b637 100644
--- a/alto/components/converter.py
+++ b/alto/components/converter.py
@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
+from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
 import torch
@@ -15,6 +16,9 @@
 
 from alto.config import Recipe
 
+if TYPE_CHECKING:
+    from torchtitan.models.base import BaseModel
+
 
 class ModelOptConverter(ModelConverter, Configurable):
 
@@ -45,6 +49,10 @@ def convert(self, model: nn.Module):
 
         for modifier in self.recipe.modifiers:
             modifier.convert(model)
+            
+    def convert_config(self, model_config: "BaseModel.Config"):
+        for modifier in self.recipe.modifiers:
+            modifier.convert_config(model_config)
 
     def pre_step(self, model_parts: list[nn.Module], **kwargs):
         for modifier in self.recipe.modifiers:
diff --git a/alto/kernels/dispatch/__init__.py b/alto/kernels/dispatch/__init__.py
index 2e03babe..fdb6beac 100644
--- a/alto/kernels/dispatch/__init__.py
+++ b/alto/kernels/dispatch/__init__.py
@@ -4,10 +4,10 @@
 
 from .config import TrainingOpConfig
 from .conversion import swap_params
-from .attention import LPScaledDotProductAttentionWrapper
+from .attention import LPScaledDotProductAttention
 
 __all__ = [
     "TrainingOpConfig",
     "swap_params",
-    "LPScaledDotProductAttentionWrapper",
+    "LPScaledDotProductAttention",
 ]
diff --git a/alto/kernels/dispatch/attention.py b/alto/kernels/dispatch/attention.py
index 33727623..b77658f2 100644
--- a/alto/kernels/dispatch/attention.py
+++ b/alto/kernels/dispatch/attention.py
@@ -3,15 +3,15 @@
 # SPDX-License-Identifier: MIT
 
 import torch
-from torchtitan.models.common.attention import (ScaledDotProductAttentionWrapper)
+from torchtitan.models.common.attention import (ScaledDotProductAttention)
 
 from alto.kernels.fp4.mxfp4.triton_flash_attention_mxfp4 import triton_attention_mxfp4
 from .config import TrainingOpConfig
 
-__all__ = ["LPScaledDotProductAttentionWrapper"]
+__all__ = ["LPScaledDotProductAttention"]
 
 
-class LPScaledDotProductAttentionWrapper(ScaledDotProductAttentionWrapper):
+class LPScaledDotProductAttention(ScaledDotProductAttention):
 
     def __init__(self, config: TrainingOpConfig):
         super().__init__()
@@ -25,17 +25,30 @@ def __init__(self, config: TrainingOpConfig):
 
     def _get_name(self) -> str:
         return f"{self.__class__.__name__}[{self.config}]"
-
+    
     def forward(
         self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
+        q_BLNH: torch.Tensor,
+        k_BLNH: torch.Tensor,
+        v_BLNH: torch.Tensor,
         *,
+        attention_masks: None = None,
         scale: float | None = None,
         enable_gqa: bool = False,
         is_causal: bool = True,
-    ):
+        **kwargs,
+    ) -> torch.Tensor:
+        if attention_masks is not None:
+            raise ValueError(
+                "ScaledDotProductAttention does not support attention_masks; it "
+                "only supports causal/non-causal attention via is_causal."
+            )
+        # Transpose to (B, N, L, H) for SDPA
+        q, k, v = (
+            q_BLNH.transpose(1, 2),
+            k_BLNH.transpose(1, 2),
+            v_BLNH.transpose(1, 2),
+        )
         batch, num_head_q, seqlen_q, head_dim_qk = q.shape
         batch_k, num_head_kv, seqlen_kv, head_dim_qk_k = k.shape
         batch_v, num_head_kv_v, seqlen_kv_v, head_dim_v = v.shape
@@ -64,4 +77,5 @@ def forward(
             use_exp2=True,
             layout="bhsd",
         )[0]
-        return o
+        # Transpose back to (B, L, N, H)
+        return o.transpose(1, 2)
diff --git a/alto/kernels/dispatch/conversion.py b/alto/kernels/dispatch/conversion.py
index 325c9e35..8b8e35b6 100644
--- a/alto/kernels/dispatch/conversion.py
+++ b/alto/kernels/dispatch/conversion.py
@@ -101,7 +101,7 @@ def post_order_traversal(
                     continue
                 module_prefix = f"{module_name}." if module_name else ""
                 full_param_name = f"{module_prefix}{cur_fqn}{'.' if cur_fqn else ''}{param_name}"
-                if target_parameter_name is None and param_name.endswith("bias"):
+                if target_parameter_name is None and "bias" in param_name:
                     logger.debug(f"Skipped {full_param_name} because it is a bias parameter")
                     continue
                 if not isinstance(param.data, TrainingWeightWrapperBaseTensor):
diff --git a/alto/kernels/dispatch/tensor.py b/alto/kernels/dispatch/tensor.py
index 40868afd..430ba838 100644
--- a/alto/kernels/dispatch/tensor.py
+++ b/alto/kernels/dispatch/tensor.py
@@ -152,6 +152,9 @@ def __tensor_unflatten__(cls, inner_tensors, flatten_spec, outer_size, outer_str
             flatten_spec["config"],
         )
 
+    def untyped_storage(self):
+        return self._data.untyped_storage()
+
     # fsdp hooks based on https://github.com/pytorch/pytorch/blob/20e40492b046b9287726d3ec656117e4dc38f0e2/test/distributed/_composable/fsdp/test_fully_shard_extensions.py#L81
     def fsdp_pre_all_gather(
         self,
diff --git a/alto/models/deepseek_v3/config_registry.py b/alto/models/deepseek_v3/config_registry.py
index dc0f0ec9..a1c7c90d 100644
--- a/alto/models/deepseek_v3/config_registry.py
+++ b/alto/models/deepseek_v3/config_registry.py
@@ -21,12 +21,12 @@
 
 def deepseek_v3_debugmodel() -> Trainer.Config:
     config = deepseek_v3_debugmodel_orig()
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 10
     config.training.local_batch_size = 4
     config.training.global_batch_size = 16
     config.training.seq_len = 2048
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -43,7 +43,7 @@ def deepseek_v3_16b() -> Trainer.Config:
     config = deepseek_v3_16b_orig()
     config.hf_assets_path = "/huggingface/hub/models--deepseek-ai--deepseek-moe-16b-base/snapshots/521d2bc4fb69a3f3ae565310fcc3b65f97af2580"
     config.dump_folder = "deepseek_v3_16b-outputs"
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.seq_len = 4096
@@ -51,7 +51,6 @@ def deepseek_v3_16b() -> Trainer.Config:
     config.metrics.enable_tensorboard = True
     config.dataloader.dataset = "c4_test"
     config.parallelism.expert_parallel_degree = 1
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
     config.checkpoint.enable = True
     config.checkpoint.initial_load_path = "/huggingface/hub/models--deepseek-ai--deepseek-moe-16b-base/snapshots/521d2bc4fb69a3f3ae565310fcc3b65f97af2580"
@@ -61,8 +60,7 @@ def deepseek_v3_16b() -> Trainer.Config:
     config.validator.dataloader.dataset = "wikitext_test"
     config.validator.freq = 10
     config.validator.steps = 10
-    config.activation_checkpoint.mode = "none"
-    config.activation_checkpoint.selective_ac_option = "1"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
diff --git a/alto/models/deepseek_v3/configs/lpt_recipe.yaml b/alto/models/deepseek_v3/configs/lpt_recipe.yaml
index cd765eef..20794b74 100644
--- a/alto/models/deepseek_v3/configs/lpt_recipe.yaml
+++ b/alto/models/deepseek_v3/configs/lpt_recipe.yaml
@@ -3,7 +3,7 @@ training_stage:
     LowPrecisionTrainingModifier:
       scheme: "mxfp4"
       targets: ["Linear", "GroupedExperts"]
-      ignore: ["output", "re:.*\\.router\\.gate"]
+      ignore: ["lm_head", "re:.*\\.router\\.gate"]
       use_2dblock_x: false
       use_2dblock_w: true
       use_hadamard: true
diff --git a/alto/models/gpt_oss/config_registry.py b/alto/models/gpt_oss/config_registry.py
index fe51292b..92232d4f 100644
--- a/alto/models/gpt_oss/config_registry.py
+++ b/alto/models/gpt_oss/config_registry.py
@@ -22,12 +22,12 @@
 
 def gpt_oss_debugmodel() -> Trainer.Config:
     config = gpt_oss_debugmodel_orig()
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 10
     config.training.local_batch_size = 4
     config.training.global_batch_size = 16
     config.training.seq_len = 2048
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -44,7 +44,7 @@ def gpt_oss_20b() -> Trainer.Config:
     config = gpt_oss_20b_orig()
     config.hf_assets_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/"
     config.dump_folder = "gpt_oss_20b-outputs"
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.seq_len = 8192
@@ -52,7 +52,6 @@ def gpt_oss_20b() -> Trainer.Config:
     config.metrics.enable_tensorboard = True
     config.dataloader.dataset = "c4_test"
     config.parallelism.expert_parallel_degree = 1
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
     config.checkpoint.enable = True
     config.checkpoint.initial_load_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/"
@@ -63,7 +62,7 @@ def gpt_oss_20b() -> Trainer.Config:
     config.validator.dataloader.dataset = "wikitext_test"
     config.validator.freq = 10
     config.validator.steps = 10
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -72,16 +71,14 @@ def gpt_oss_20b_pretrain() -> Trainer.Config:
     config = gpt_oss_20b_orig()
     config.hf_assets_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/"
     config.dump_folder = "gpt_oss_20b-pretrain-subset-lr4e-4-outputs"
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 1200000
     config.training.local_batch_size = 1
     config.training.global_batch_size = 16
     config.training.seq_len = 8192
-    config.optimizer.lr = 4e-4
-    config.optimizer.weight_decay = 0.1
-    config.optimizer.beta1 = 0.9
-    config.optimizer.beta2 = 0.95
-    config.optimizer.eps = 1e-5
+    config.optimizer.param_groups[0].optimizer_kwargs["lr"] = 4e-4
+    config.optimizer.param_groups[0].optimizer_kwargs["betas"] = (0.9, 0.95)
+    config.optimizer.param_groups[0].optimizer_kwargs["eps"] = 1e-5
     config.lr_scheduler.min_lr_factor = 0.1
     config.lr_scheduler.warmup_steps = 128
     config.lr_scheduler.decay_ratio = 1 - 128 / config.training.steps
@@ -91,7 +88,6 @@ def gpt_oss_20b_pretrain() -> Trainer.Config:
     config.dataloader.dataset = "megatron"
     config.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-train.en_6_text_document.idx"
     config.parallelism.expert_parallel_degree = 8
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
     config.checkpoint.enable = True
     config.checkpoint.interval = 1000
@@ -101,14 +97,14 @@ def gpt_oss_20b_pretrain() -> Trainer.Config:
     config.validator.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-validation-91205-samples.en_text_document.idx"
     config.validator.freq = 768
     config.validator.steps = 64
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
 
 def gpt_oss_20b_lpt() -> Trainer.Config:
     config = gpt_oss_20b_pretrain()
-    config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-outputs"
+    config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-refactor-outputs"
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe.yaml",),
     ],)
diff --git a/alto/models/gpt_oss/configs/lpt_recipe.yaml b/alto/models/gpt_oss/configs/lpt_recipe.yaml
index ad935b53..f2d78836 100644
--- a/alto/models/gpt_oss/configs/lpt_recipe.yaml
+++ b/alto/models/gpt_oss/configs/lpt_recipe.yaml
@@ -4,7 +4,7 @@ training_stage:
       scheme: "mxfp4"
       targets: ["Linear", "GptOssGroupedExperts"]
       # targets: ["Linear"]
-      ignore: ["output", "re:.*\\.router\\.gate"]
+      ignore: ["lm_head", "re:.*\\.router\\.gate"]
       use_2dblock_x: false
       use_2dblock_w: true
       use_hadamard: true
diff --git a/alto/models/llama3/config_registry.py b/alto/models/llama3/config_registry.py
index 097d1256..1d9acf2a 100644
--- a/alto/models/llama3/config_registry.py
+++ b/alto/models/llama3/config_registry.py
@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
-from torchtitan.components.optimizer import OptimizersContainer
+from torchtitan.components.optimizer import default_adamw
 from torchtitan.trainer import Trainer
 from torchtitan.protocols.model_converter import ModelConvertersContainer
 from torchtitan.models.llama3.config_registry import (
@@ -54,12 +54,12 @@
 
 def llama3_debugmodel() -> Trainer.Config:
     config = llama3_debugmodel_orig()
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 4
     config.training.global_batch_size = 16
     config.training.seq_len = 2048
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -85,13 +85,13 @@ def llama3_1b() -> Trainer.Config:
     config = llama3_1b_orig()
     config.hf_assets_path = "/group/archive_dataset_6_nobkup/archive_modelzoo/sequence_learning/weights/nlp-pretrained-model/meta-llama/Llama-3.2-1B"
     config.metrics.log_freq = 1
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.global_batch_size = 10
     config.training.seq_len = 8192
     config.dataloader.dataset = "c4_test"
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.checkpoint.enable = True
     config.checkpoint.interval = 10
     config.checkpoint.initial_load_path = "/group/archive_dataset_6_nobkup/archive_modelzoo/sequence_learning/weights/nlp-pretrained-model/meta-llama/Llama-3.2-1B"
@@ -107,7 +107,7 @@ def llama3_1b() -> Trainer.Config:
 def llama3_1b_opt() -> Trainer.Config:
     config = llama3_1b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/recipe.yaml",),
     ],)
@@ -128,12 +128,12 @@ def llama3_8b_pretrain() -> Trainer.Config:
     config.dump_folder = "llama3_8b-mi308-pretrain-subset-gbs384-lr1e-4-outputs"
     config.metrics.log_freq = 1
     config.metrics.enable_tensorboard = True
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 5000
     config.training.local_batch_size = 2
     config.training.global_batch_size = 384
     config.training.seq_len = 8192
-    config.optimizer.lr = 1e-4
+    config.optimizer.param_groups[0].optimizer_kwargs["lr"] = 1e-4
     config.lr_scheduler.min_lr_factor = 0.0
     config.lr_scheduler.warmup_steps = 500
     config.lr_scheduler.decay_ratio = 0.9
@@ -141,9 +141,8 @@ def llama3_8b_pretrain() -> Trainer.Config:
     config.dataloader.dataset = "megatron"
     config.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-train.en_6_text_document.idx"
     config.parallelism.expert_parallel_degree = 1
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.checkpoint.enable = False
     config.checkpoint.interval = 10
     config.checkpoint.initial_load_path = "/huggingface/hub/models--unsloth--Llama-3.1-8B/snapshots/3f0d51f8e5640f98f1a96ea9044a0e55c0a83814"
@@ -177,7 +176,7 @@ def llama3_8b_lpt() -> Trainer.Config:
 def llama3_1b_gptq() -> Trainer.Config:
     config = llama3_1b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/gptq_recipe.yaml",),
     ],)
@@ -187,7 +186,7 @@ def llama3_1b_gptq() -> Trainer.Config:
 def llama3_1b_awq() -> Trainer.Config:
     config = llama3_1b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/awq_recipe.yaml",),
     ],)
@@ -197,7 +196,7 @@ def llama3_1b_awq() -> Trainer.Config:
 def llama3_1b_mx9_wa() -> Trainer.Config:
     config = llama3_1b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/mx9_wa_recipe.yaml",),
     ],)
@@ -207,7 +206,7 @@ def llama3_1b_mx9_wa() -> Trainer.Config:
 def llama3_1b_mx6_wa() -> Trainer.Config:
     config = llama3_1b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/mx6_wa_recipe.yaml",),
     ],)
@@ -221,13 +220,13 @@ def llama3_8b() -> Trainer.Config:
     config = llama3_8b_orig()
     config.hf_assets_path = LLAMA3_8B_PATH
     config.metrics.log_freq = 1
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.global_batch_size = 8
     config.training.seq_len = 2048
     config.dataloader = HuggingFaceTextDataLoader.Config(dataset="c4_test")
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.checkpoint.enable = True
     config.checkpoint.interval = 10
     config.checkpoint.initial_load_path = LLAMA3_8B_PATH
@@ -243,7 +242,7 @@ def llama3_8b() -> Trainer.Config:
 def llama3_8b_gptq() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.training.global_batch_size = 128
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/gptq_recipe.yaml",),
@@ -254,7 +253,7 @@ def llama3_8b_gptq() -> Trainer.Config:
 def llama3_8b_rtn() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/rtn_recipe.yaml",),
     ],)
@@ -264,7 +263,7 @@ def llama3_8b_rtn() -> Trainer.Config:
 def llama3_8b_awq() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/awq_recipe.yaml",),
     ],)
@@ -279,7 +278,7 @@ def llama3_8b_awq() -> Trainer.Config:
 def llama3_8b_wanda() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/wanda_recipe.yaml",),
     ],)
@@ -289,7 +288,7 @@ def llama3_8b_wanda() -> Trainer.Config:
 def llama3_8b_sparsegpt() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.training.global_batch_size = 128
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/sparsegpt_recipe.yaml",),
@@ -300,7 +299,7 @@ def llama3_8b_sparsegpt() -> Trainer.Config:
 def llama3_8b_magnitude() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/magnitude_recipe.yaml",),
     ],)
@@ -310,7 +309,7 @@ def llama3_8b_magnitude() -> Trainer.Config:
 def llama3_8b_admm() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.training.global_batch_size = 128
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/admm_recipe.yaml",),
@@ -321,7 +320,7 @@ def llama3_8b_admm() -> Trainer.Config:
 def llama3_8b_alps() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.training.global_batch_size = 128
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/alps_recipe.yaml",),
@@ -337,7 +336,7 @@ def llama3_8b_alps() -> Trainer.Config:
 def llama3_8b_wanda_structured() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/wanda_structured_recipe.yaml",),
     ],)
@@ -347,7 +346,7 @@ def llama3_8b_wanda_structured() -> Trainer.Config:
 def llama3_8b_obs() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.training.global_batch_size = 128
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/obs_recipe.yaml",),
@@ -358,7 +357,7 @@ def llama3_8b_obs() -> Trainer.Config:
 def llama3_8b_admm_structured() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.training.global_batch_size = 128
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/admm_structured_recipe.yaml",),
@@ -369,7 +368,7 @@ def llama3_8b_admm_structured() -> Trainer.Config:
 def llama3_8b_cosine_similarity() -> Trainer.Config:
     config = llama3_8b()
     config.training.steps = 1
-    config.optimizer = OptimizersContainer.Config(lr=0.0)
+    config.optimizer = default_adamw(lr=0.0)
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/llama3/configs/cosine_similarity_recipe.yaml",),
     ],)
@@ -385,13 +384,13 @@ def instella_3b() -> Trainer.Config:
     config = instella_3b_orig()
     config.hf_assets_path = "/group/ossmodelzoo/hanwang2/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e"
     config.metrics.log_freq = 1
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.global_batch_size = 10
     config.training.seq_len = 4096
     config.dataloader.dataset = "c4_test"
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.checkpoint.enable = True
     config.checkpoint.interval = 10
     config.checkpoint.initial_load_path = "/group/ossmodelzoo/hanwang2/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e"
diff --git a/alto/models/llama3/configs/admm_recipe.yaml b/alto/models/llama3/configs/admm_recipe.yaml
index b9bdd6f8..416b7a92 100644
--- a/alto/models/llama3/configs/admm_recipe.yaml
+++ b/alto/models/llama3/configs/admm_recipe.yaml
@@ -11,4 +11,4 @@ sparsification_stage:
       admm_pruning_granularity: 32
       admm_iterations: 32
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/admm_structured_recipe.yaml b/alto/models/llama3/configs/admm_structured_recipe.yaml
index 17501d64..a1a381ac 100644
--- a/alto/models/llama3/configs/admm_structured_recipe.yaml
+++ b/alto/models/llama3/configs/admm_structured_recipe.yaml
@@ -10,4 +10,4 @@ pruning_stage:
       admm_iterations: 16
       dampening_frac: 0.01
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/alps_recipe.yaml b/alto/models/llama3/configs/alps_recipe.yaml
index c2530745..633b1600 100644
--- a/alto/models/llama3/configs/alps_recipe.yaml
+++ b/alto/models/llama3/configs/alps_recipe.yaml
@@ -12,4 +12,4 @@ sparsification_stage:
       alps_update_iter: 16
       alps_switch_iter: 16
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/awq_debug_recipe.yaml b/alto/models/llama3/configs/awq_debug_recipe.yaml
index b329af00..4bc21c91 100644
--- a/alto/models/llama3/configs/awq_debug_recipe.yaml
+++ b/alto/models/llama3/configs/awq_debug_recipe.yaml
@@ -5,7 +5,7 @@ quantization_stage:
     AWQModifier:
       sequential_granularity: "sublayer"
       n_grid: 20
-      ignore: ["output"]
+      ignore: ["lm_head"]
       mappings:
         - smooth_layer: "re:.*attention_norm"
           balance_layers:
diff --git a/alto/models/llama3/configs/awq_recipe.yaml b/alto/models/llama3/configs/awq_recipe.yaml
index d33c039a..0ab88397 100644
--- a/alto/models/llama3/configs/awq_recipe.yaml
+++ b/alto/models/llama3/configs/awq_recipe.yaml
@@ -6,7 +6,7 @@ quantization_stage:
     AWQModifier:
       sequential_granularity: "sublayer"
       n_grid: 20
-      ignore: ["output"]
+      ignore: ["lm_head"]
       mappings:
         - smooth_layer: "re:.*attention_norm"
           balance_layers:
diff --git a/alto/models/llama3/configs/cosine_similarity_recipe.yaml b/alto/models/llama3/configs/cosine_similarity_recipe.yaml
index e9b4d399..827f7d3c 100644
--- a/alto/models/llama3/configs/cosine_similarity_recipe.yaml
+++ b/alto/models/llama3/configs/cosine_similarity_recipe.yaml
@@ -8,4 +8,4 @@ pruning_stage:
       pruning_dimension: "layer"
       targets:
         - "re:.*\\.wq$"
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/gptq_recipe.yaml b/alto/models/llama3/configs/gptq_recipe.yaml
index 11721ead..5d162681 100644
--- a/alto/models/llama3/configs/gptq_recipe.yaml
+++ b/alto/models/llama3/configs/gptq_recipe.yaml
@@ -7,7 +7,7 @@ quantization_stage:
       sequential_granularity: "sublayer"
       block_size: 128
       dampening_frac: 0.01
-      ignore: ["output"]
+      ignore: ["lm_head"]
       config_groups:
         group_0:
           weights:
diff --git a/alto/models/llama3/configs/guanchen_recipe.yaml b/alto/models/llama3/configs/guanchen_recipe.yaml
index 6c25fc9b..f31232ae 100644
--- a/alto/models/llama3/configs/guanchen_recipe.yaml
+++ b/alto/models/llama3/configs/guanchen_recipe.yaml
@@ -4,7 +4,7 @@
 #       sparsity: 0.5
 #       mask_structure: "2:4"
 #       targets: ["Linear"]
-#       ignore: ["output"]
+#       ignore: ["lm_head"]
 
 # sparsity_stage:
 #   sparsity_modifiers:
@@ -12,7 +12,7 @@
 #       sparsity: 0.5
 #       mask_structure: "2:4"
 #       targets: ["Linear"]
-#       ignore: ["output"]
+#       ignore: ["lm_head"]
 
 # sparsity_stage:
 #   sparsity_modifiers:
@@ -20,7 +20,7 @@
 #       sparsity: 0.5
 #       mask_structure: "2:4"
 #       targets: ["Linear"]
-#       ignore: ["output"]
+#       ignore: ["lm_head"]
 
 # sparsity_stage:
 #   sparsity_modifiers:
@@ -28,7 +28,7 @@
 #       sparsity: 0.5
 #       mask_structure: "2:4"
 #       targets: ["Linear"]
-#       ignore: ["output"]
+#       ignore: ["lm_head"]
 
 pruning_stage:
   pruning_modifiers:
diff --git a/alto/models/llama3/configs/lpt_recipe.yaml b/alto/models/llama3/configs/lpt_recipe.yaml
index 2a12a1d1..3a6a36cd 100644
--- a/alto/models/llama3/configs/lpt_recipe.yaml
+++ b/alto/models/llama3/configs/lpt_recipe.yaml
@@ -3,7 +3,7 @@ training_stage:
     LowPrecisionTrainingModifier:
       scheme: "mxfp4"
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
       use_2dblock_x: false
       use_2dblock_w: true
       use_hadamard: true
diff --git a/alto/models/llama3/configs/magnitude_recipe.yaml b/alto/models/llama3/configs/magnitude_recipe.yaml
index c9d4bec3..ed3ccef7 100644
--- a/alto/models/llama3/configs/magnitude_recipe.yaml
+++ b/alto/models/llama3/configs/magnitude_recipe.yaml
@@ -7,4 +7,4 @@ sparsification_stage:
       sparsity: 0.5
       mask_structure: "0:0"
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/mx6_wa_recipe.yaml b/alto/models/llama3/configs/mx6_wa_recipe.yaml
index 5ab4d3b1..d04d6d13 100644
--- a/alto/models/llama3/configs/mx6_wa_recipe.yaml
+++ b/alto/models/llama3/configs/mx6_wa_recipe.yaml
@@ -4,7 +4,7 @@
 quantization_stage:
   quantization_modifiers:
     QuantizationModifier:
-      ignore: ["output"]          # lm_head not quantized
+      ignore: ["lm_head"]          # lm_head not quantized
       sequential: false
       config_groups:
         group_0:
diff --git a/alto/models/llama3/configs/mx9_wa_recipe.yaml b/alto/models/llama3/configs/mx9_wa_recipe.yaml
index 3df48414..93ea42b9 100644
--- a/alto/models/llama3/configs/mx9_wa_recipe.yaml
+++ b/alto/models/llama3/configs/mx9_wa_recipe.yaml
@@ -4,7 +4,7 @@
 quantization_stage:
   quantization_modifiers:
     QuantizationModifier:
-      ignore: ["output"]          # lm_head not quantized
+      ignore: ["lm_head"]          # lm_head not quantized
       sequential: false
       config_groups:
         group_0:
diff --git a/alto/models/llama3/configs/obs_recipe.yaml b/alto/models/llama3/configs/obs_recipe.yaml
index 2447c806..0b214a9b 100644
--- a/alto/models/llama3/configs/obs_recipe.yaml
+++ b/alto/models/llama3/configs/obs_recipe.yaml
@@ -10,4 +10,4 @@ pruning_stage:
       attention_pruning_granularity: 1
       dampening_frac: 0.01
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/recipe.yaml b/alto/models/llama3/configs/recipe.yaml
index bc4da067..50745533 100644
--- a/alto/models/llama3/configs/recipe.yaml
+++ b/alto/models/llama3/configs/recipe.yaml
@@ -4,11 +4,11 @@
 #       sparsity: 0.5
 #       mask_structure: "2:4"
 #       targets: ["Linear"]
-#       ignore: ["output"]
+#       ignore: ["lm_head"]
 quantization_stage:
   quantization_modifiers:
     QuantizationModifier:
-      ignore: ["output"]
+      ignore: ["lm_head"]
       config_groups:
         group_0:
           weights:
diff --git a/alto/models/llama3/configs/rtn_recipe.yaml b/alto/models/llama3/configs/rtn_recipe.yaml
index 072ec777..8bfc334e 100644
--- a/alto/models/llama3/configs/rtn_recipe.yaml
+++ b/alto/models/llama3/configs/rtn_recipe.yaml
@@ -1,7 +1,7 @@
 quantization_stage:
   quantization_modifiers:
     QuantizationModifier:
-      ignore: ["output"]
+      ignore: ["lm_head"]
       config_groups:
         group_0:
           weights:
diff --git a/alto/models/llama3/configs/sparsegpt_recipe.yaml b/alto/models/llama3/configs/sparsegpt_recipe.yaml
index f785b91c..e2785804 100644
--- a/alto/models/llama3/configs/sparsegpt_recipe.yaml
+++ b/alto/models/llama3/configs/sparsegpt_recipe.yaml
@@ -9,4 +9,4 @@ sparsification_stage:
       block_size: 128
       dampening_frac: 0.01
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/wanda_recipe.yaml b/alto/models/llama3/configs/wanda_recipe.yaml
index 0f3aa082..de94fe02 100644
--- a/alto/models/llama3/configs/wanda_recipe.yaml
+++ b/alto/models/llama3/configs/wanda_recipe.yaml
@@ -7,4 +7,4 @@ sparsification_stage:
       sparsity: 0.5
       mask_structure: "0:0"
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/llama3/configs/wanda_structured_recipe.yaml b/alto/models/llama3/configs/wanda_structured_recipe.yaml
index b5541d9c..57791005 100644
--- a/alto/models/llama3/configs/wanda_structured_recipe.yaml
+++ b/alto/models/llama3/configs/wanda_structured_recipe.yaml
@@ -7,4 +7,4 @@ pruning_stage:
       sparsity: 0.25
       pruning_dimension: "mlp"
       targets: ["Linear"]
-      ignore: ["output"]
+      ignore: ["lm_head"]
diff --git a/alto/models/patcher.py b/alto/models/patcher.py
index 8c751354..3059b0f4 100644
--- a/alto/models/patcher.py
+++ b/alto/models/patcher.py
@@ -101,14 +101,14 @@ def fake_quantize(x, scale, zero_point, args, g_idx, global_scale):
 
     @classmethod
     def patch_apply_rotary_emb_complex(cls):
-        from torchtitan.models.common import rope, attention
-        original_apply_rotary_emb_complex = rope.apply_rotary_emb_complex
+        from torchtitan.models.common.rope import ComplexRoPE
+        original_apply_rotary_emb_complex = ComplexRoPE.apply_rotary_emb
 
         def apply_rotary_emb_complex(
+            cls,
             xq: torch.Tensor,
             xk: torch.Tensor,
-            freqs_cis: torch.Tensor,
-            positions: torch.Tensor | None = None,
+            rope_cache: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             head_dim = xq.shape[-1]
             xq = xq.reshape(
@@ -127,8 +127,6 @@ def apply_rotary_emb_complex(
                 *xk.shape[:-1],
                 head_dim,
             ).contiguous()
-            return original_apply_rotary_emb_complex(xq, xk, freqs_cis, positions)
+            return original_apply_rotary_emb_complex(xq, xk, rope_cache)
 
-        rope.apply_rotary_emb_complex = apply_rotary_emb_complex
-        attention.apply_rotary_emb_complex = apply_rotary_emb_complex
-        patch("torchtitan.models.common.rope.apply_rotary_emb_complex", apply_rotary_emb_complex).__enter__()
+        ComplexRoPE.apply_rotary_emb = classmethod(apply_rotary_emb_complex)
diff --git a/alto/modifiers/base.py b/alto/modifiers/base.py
index b6b71697..33c190a1 100644
--- a/alto/modifiers/base.py
+++ b/alto/modifiers/base.py
@@ -9,12 +9,15 @@
 
 # modified from https://github.com/vllm-project/llm-compressor/blob/f3f14af3ee56e35db7e1faf6da8833f84a570baf/src/llmcompressor/modifiers/modifier.py
 
+from typing import TYPE_CHECKING
 from abc import abstractmethod
 from pydantic import ConfigDict
 import torch
 from torch.nn import Module
 from torchtitan.tools.logging import logger
 from alto.modifiers.utils import HooksMixin
+if TYPE_CHECKING:
+    from torchtitan.protocols.model import BaseModel
 
 __all__ = ["Modifier"]
 
@@ -52,6 +55,9 @@ def requires_training_mode(self) -> bool:
 
     def convert(self, model: Module, **kwargs) -> bool:
         return self.on_convert(model, **kwargs)
+    
+    def convert_config(self, model_config: "BaseModel.Config") -> bool:
+        return self.on_convert_config(model_config)
 
     def initialize(self, model_parts: list[Module], **kwargs):
         if self.initialized_:
@@ -102,6 +108,10 @@ def on_post_step(self, model_parts: list[Module], **kwargs) -> bool:
     def on_convert(self, model: Module, **kwargs) -> bool:
         raise NotImplementedError
 
+    @abstractmethod
+    def on_convert_config(self, model_config: "BaseModel.Config") -> bool:
+        raise NotImplementedError
+
     def _infer_sequential_targets(self, model: torch.nn.Module) -> str | list[str]:
         match self.sequential_targets:
             case None:
diff --git a/alto/modifiers/distillation/base.py b/alto/modifiers/distillation/base.py
index 803cec01..ec7a7a5b 100644
--- a/alto/modifiers/distillation/base.py
+++ b/alto/modifiers/distillation/base.py
@@ -4,7 +4,7 @@
 
 from abc import abstractmethod
 from functools import partial
-from typing import Any, Literal
+from typing import Any, Literal, TYPE_CHECKING
 import gc
 
 import torch
@@ -14,7 +14,7 @@
 from torchtitan.tools.logging import logger
 from torchtitan.tools.utils import device_type
 from torchtitan.components.loss import IGNORE_INDEX
-from torchtitan.components.optimizer import OptimizersContainer
+from torchtitan.components.optimizer import OptimizersContainer, ParamGroupConfig
 from torchtitan.components.lr_scheduler import LRSchedulersContainer
 
 from alto.observers import Observer
@@ -23,6 +23,8 @@
 from alto.utils.pytorch.module import (
     get_layers,)
 from alto.modifiers.distillation.utils import losses
+if TYPE_CHECKING:
+    from torchtitan.protocols.model import BaseModel
 
 TEACHER_OBSERVER_BASE_NAME = "output"
 STUDENT_OBSERVER_BASE_NAME = "student_output"
@@ -271,12 +273,18 @@ def _get_loss_fn(self, name: str):
 
     def _build_optimizers(self, model_parts: list[Module]):
         config = OptimizersContainer.Config(
-            name=self.optimizer,
-            lr=self.lr,
-            beta1=self.beta1,
-            beta2=self.beta2,
-            eps=self.eps,
-            weight_decay=self.weight_decay,
+            param_groups=[
+                ParamGroupConfig(
+                    pattern=r".*",
+                    optimizer_name=self.optimizer,
+                    optimizer_kwargs={
+                        "lr": self.lr,
+                        "betas": (self.beta1, self.beta2),
+                        "eps": self.eps,
+                        "weight_decay": self.weight_decay,
+                    },
+                )
+            ],
             implementation=self.implementation,
         )
 
@@ -284,3 +292,6 @@ def _build_optimizers(self, model_parts: list[Module]):
 
     def on_convert(self, model: Module, **kwargs) -> bool:
         return True
+    
+    def on_convert_config(self, model_config: "BaseModel.Config") -> bool:
+        return True
diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py
index 319d0835..43e2f941 100644
--- a/alto/modifiers/lpt/base.py
+++ b/alto/modifiers/lpt/base.py
@@ -2,25 +2,27 @@
 #
 # SPDX-License-Identifier: MIT
 
-from typing import Literal
+from typing import Literal, Iterable, TYPE_CHECKING
 import torch
 from torch.nn import Module
-from compressed_tensors.utils import match_named_modules
+from compressed_tensors.utils import match_named_modules, match_name
 from pydantic import PrivateAttr, Field, field_validator, model_validator
-from torchtitan.models.common.attention import BaseAttention
-from torchtitan.models.common.moe.utils import set_token_group_alignment_size_m
+from torchtitan.models.common.attention import BaseAttention, ScaledDotProductAttention
 from torchtitan.tools.logging import logger
 
 from alto.modifiers import Modifier
 from alto.kernels.dispatch import (
     swap_params,
     TrainingOpConfig,
-    LPScaledDotProductAttentionWrapper,
+    LPScaledDotProductAttention,
 )
 from alto.kernels.fp4.mxfp4.mxfp_grouped_gemm.autotune import ALIGN_SIZE_M
 from alto.nn import DecomposedLinear
 from alto.components.optimizer import DeOscillationConfig, enable_de_oscillation
 
+if TYPE_CHECKING:
+    from torchtitan.protocols.model import BaseModel
+
 __all__ = ["LowPrecisionTrainingModifier"]
 
 
@@ -71,8 +73,6 @@ class LowPrecisionTrainingModifier(Modifier):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
-        set_token_group_alignment_size_m(ALIGN_SIZE_M)
-
     @field_validator("targets", mode="before")
     def validate_targets(cls, value: str | list[str]) -> list[str]:
         if isinstance(value, str):
@@ -166,8 +166,8 @@ def on_convert(self, model: Module, **kwargs) -> bool:
         for scheme_obj, targets in self.resolved_config.items():
             for name, module in match_named_modules(model, targets, self.ignore):
                 if isinstance(module, BaseAttention):
-                    assert module.attn_backend == "sdpa", "Only SDPA attention is supported for now."
-                    module.inner_attention = LPScaledDotProductAttentionWrapper(config=scheme_obj)
+                    assert isinstance(module.inner_attention, ScaledDotProductAttention), "Only SDPA attention is supported for now."
+                    module.inner_attention = LPScaledDotProductAttention(config=scheme_obj)
                 elif isinstance(module, torch.nn.Linear):
                     if self.lora_rank > 0:
                         module = DecomposedLinear.from_linear(module, lora_rank=self.lora_rank)
@@ -185,11 +185,48 @@ def on_convert(self, model: Module, **kwargs) -> bool:
         logger.info(f"LowPrecisionTrainingModifier converted model: {model}")
         return True
 
+    def on_convert_config(self, model_config: "BaseModel.Config") -> bool:
+        # convert configs to enable token alignment
+        from torchtitan.models.common.moe import GroupedExperts
+        from torchtitan.components.quantization.utils import swap_token_dispatcher
+        from torchtitan.models.common.nn_modules import Linear
+
+        for _fqn, config, parent, attr in model_config.traverse(GroupedExperts.Config):
+            swap_token_dispatcher(config, ALIGN_SIZE_M)
+
+        def is_match(
+            name: str,
+            module_cls_name: str,
+            targets: str | Iterable[str],
+            ignore: str | Iterable[str] = tuple(),
+        ) -> bool:
+            targets = [targets] if isinstance(targets, str) else targets
+            ignore = [ignore] if isinstance(ignore, str) else ignore
+
+            return any(
+                match_name(name, target) or module_cls_name == target
+                for target in targets
+            ) and not any(
+                match_name(name, ign) or module_cls_name == ign for ign in ignore
+            )
+
+        # replace Linear with DecomposedLinear
+        if self.lora_rank > 0:
+            resolved_targets = list(self.resolved_config.values())
+            for _fqn, config, parent, attr in model_config.traverse(Linear.Config):
+                for target in resolved_targets:
+                    if is_match(_fqn, "Linear", target, self.ignore):
+                        new_config = DecomposedLinear.Config(
+                            in_features=config.in_features,
+                            out_features=config.out_features,
+                            bias=config.bias,
+                            lora_rank=self.lora_rank,
+                            param_init=config.param_init | DecomposedLinear._EXTRA_INIT,
+                        )
+                        setattr(parent, attr, new_config)
+        return True
+
     def on_initialize(self, model_parts: list[Module], **kwargs) -> bool:
-        for model_part in model_parts:
-            for child in model_part.modules():
-                if isinstance(child, DecomposedLinear):
-                    child.init_lora_weights(init_std=0.02)
         return True
 
     def on_pre_step(self, model_parts: list[Module], **kwargs) -> bool:
diff --git a/alto/modifiers/pruning/base.py b/alto/modifiers/pruning/base.py
index 9f229e04..fb965bb3 100644
--- a/alto/modifiers/pruning/base.py
+++ b/alto/modifiers/pruning/base.py
@@ -11,7 +11,7 @@
 
 from abc import abstractmethod
 from functools import partial
-from typing import Any, Generator, Literal
+from typing import Any, Generator, Literal, TYPE_CHECKING
 import gc
 import re
 
@@ -30,6 +30,8 @@
     get_prunable_layers,
     match_targets,
 )
+if TYPE_CHECKING:
+    from torchtitan.protocols.model import BaseModel
 
 LAYER_OBSERVER_BASE_NAME = "pruning"
 PruningDim = Literal["attn", "mlp", "attn+mlp", "mlp+attn", "layer", "sublayer", "hidden_dim"]
@@ -233,6 +235,9 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool:
     def on_convert(self, model: Module, **kwargs) -> bool:
         return True
 
+    def on_convert_config(self, model_config: "BaseModel.Config") -> bool:
+        return True
+
     # ---- sequential helpers -------------------------------------------
 
     def _capture_hook(self, module: Module, args: Any, kwargs: Any = None):
diff --git a/alto/modifiers/quantization/base.py b/alto/modifiers/quantization/base.py
index 3511fdd2..185f3217 100644
--- a/alto/modifiers/quantization/base.py
+++ b/alto/modifiers/quantization/base.py
@@ -11,7 +11,7 @@
 # Subclasses override _process_block() for algorithm-specific per-block logic.
 
 import gc
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union, TYPE_CHECKING
 
 import torch
 import tqdm
@@ -29,6 +29,9 @@
 from alto.modifiers.quantization.calibration import update_weight_zp_scale
 from alto.utils.pytorch.module import get_layers
 
+if TYPE_CHECKING:
+    from torchtitan.protocols.model import BaseModel
+
 __all__ = ["QuantizationModifier"]
 
 
@@ -70,11 +73,6 @@ class QuantizationModifier(Modifier, QuantizationMixin):
     # ---- lifecycle ----------------------------------------------------
 
     def on_initialize(self, model_parts: list[Module], **kwargs) -> bool:
-        if not QuantizationMixin.has_config(self):
-            raise ValueError("QuantizationModifier requires that quantization fields be specified")
-        for m in model_parts:
-            QuantizationMixin.initialize_quantization(self, m)
-
         if self.sequential:
             self._build_sequential_blocks(model_parts)
         return True
@@ -110,6 +108,25 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool:
         return True
 
     def on_convert(self, model: Module, **kwargs) -> bool:
+        if not QuantizationMixin.has_config(self):
+            raise ValueError("QuantizationModifier requires that quantization fields be specified")
+        # Note: qparams will be registered in the initialize_quantization method,
+        #       so it has to be done before applying parallelism
+        QuantizationMixin.initialize_quantization(self, model)
+        
+        # patch param_init dict because the qparams are registered on meta device
+        # and need to be initialized after to_empty copy.
+        for mod_name, mod in model.named_modules():
+            qscheme = getattr(mod, "quantization_scheme", None)
+            if qscheme is not None:
+                for pname, p in mod.named_parameters():
+                    if pname.endswith("scale"):
+                        mod._param_init[pname] = torch.nn.init.ones_
+                    elif pname.endswith("zero_point"):
+                        mod._param_init[pname] = torch.nn.init.zeros_
+        return True
+
+    def on_convert_config(self, model_config: "BaseModel.Config") -> bool:
         return True
 
     # ---- sequential loop (template) -----------------------------------
diff --git a/alto/modifiers/sparsification/base.py b/alto/modifiers/sparsification/base.py
index 2e3d94ff..6d667b1a 100644
--- a/alto/modifiers/sparsification/base.py
+++ b/alto/modifiers/sparsification/base.py
@@ -11,7 +11,7 @@
 
 from abc import abstractmethod
 from functools import partial
-from typing import Any, Generator, Literal
+from typing import Any, Generator, Literal, TYPE_CHECKING
 import gc
 import re
 
@@ -30,6 +30,8 @@
     get_prunable_layers,
     match_targets,
 )
+if TYPE_CHECKING:
+    from torchtitan.protocols.model import BaseModel
 
 LAYER_OBSERVER_BASE_NAME = "sparsity"
 
@@ -242,6 +244,9 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool:
     def on_convert(self, model: Module, **kwargs) -> bool:
         return True
 
+    def on_convert_config(self, model_config: "BaseModel.Config") -> bool:
+        return True
+
     # ---- sequential helpers -------------------------------------------
 
     def _capture_hook(self, module: Module, args: Any, kwargs: Any = None):
diff --git a/alto/modifiers/sparsification/wanda.py b/alto/modifiers/sparsification/wanda.py
index 24c62df2..10abf04c 100644
--- a/alto/modifiers/sparsification/wanda.py
+++ b/alto/modifiers/sparsification/wanda.py
@@ -11,6 +11,7 @@
 
 import torch
 from torch.nn import Module
+from torch.distributed.tensor import DTensor, Replicate
 from torchtitan.tools.logging import logger
 
 from alto.modifiers.sparsification.base import SparsityModifierBase
@@ -79,6 +80,11 @@ def compress_modules(self):
                 prune_n=self._prune_n,
                 prune_m=self._prune_m,
             )
+            if isinstance(module.weight, DTensor):
+                sparsified_weight = sparsified_weight.redistribute(
+                    module.weight.device_mesh,
+                    placements=module.weight.placements,
+                )
             module.weight.data.copy_(sparsified_weight)
 
     def _sparsify_weight(
@@ -111,6 +117,12 @@ def _sparsify_weight(
 
         W = W.to(dtype=PRECISION)
         S = row_scalar
+        
+        device_mesh = None
+        if isinstance(W, DTensor):
+            device_mesh = W.device_mesh
+            W = W.redistribute(device_mesh, placements=(Replicate(),)).to_local()
+            S = S.redistribute(device_mesh, placements=(Replicate(),)).to_local()
 
         W_metric = torch.abs(W) * torch.sqrt(S.reshape((1, -1)))
 
@@ -138,4 +150,8 @@ def _sparsify_weight(
 
         W = W.reshape(final_shape).to(final_dtype)
 
+        if device_mesh is not None:
+            W = DTensor.from_local(W, device_mesh=device_mesh, placements=(Replicate(),))
+            W_mask = DTensor.from_local(W_mask, device_mesh=device_mesh, placements=(Replicate(),))
+
         return W, W_mask.reshape(final_shape)
diff --git a/alto/nn/decomposed_linear.py b/alto/nn/decomposed_linear.py
index 54e721a2..c5e5bfaf 100644
--- a/alto/nn/decomposed_linear.py
+++ b/alto/nn/decomposed_linear.py
@@ -1,23 +1,66 @@
 # Copyright (c) 2026 Advanced Micro Devices, Inc.
 #
 # SPDX-License-Identifier: MIT
-
+from dataclasses import dataclass
+from functools import partial
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from torchtitan.protocols.module import Module
+
+class DecomposedLinear(Module):
+    
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        in_features: int
+        out_features: int
+        bias: bool = False
+        lora_rank: int = 32
+        
+    _EXTRA_INIT = {
+        "u": partial(nn.init.normal_, std=0.02),
+        "v": nn.init.zeros_,
+        "sigma": nn.init.ones_,
+    }
+
+    def __init__(self, config: Config):
+        super().__init__()
+        self.in_features = config.in_features
+        self.out_features = config.out_features
+
+        self.weight = nn.Parameter(torch.empty(self.out_features, self.in_features))
+        self.bias = nn.Parameter(torch.empty(self.out_features)) if config.bias else None
+        self.u = nn.Parameter(torch.empty(config.lora_rank, self.out_features))  # transposed
+        self.v = nn.Parameter(torch.empty(self.in_features, config.lora_rank))
+        self.sigma = nn.Parameter(torch.empty(config.lora_rank))
 
-class DecomposedLinear(nn.Module):
-    def __init__(self, in_features, out_features, bias=True, lora_rank=32):
-        super(DecomposedLinear, self).__init__()
-        self.in_features = in_features
-        self.out_features = out_features
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, lora_rank: int = 32) -> "DecomposedLinear":
+        config = cls.Config(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            bias=linear.bias is not None,
+            lora_rank=lora_rank,
+        )
+        module = cls(config).to(device=linear.weight.device, dtype=linear.weight.dtype)
 
-        self.weight = nn.Parameter(torch.empty(out_features, in_features))
-        self.bias = nn.Parameter(torch.empty(out_features)) if bias else None
-        self.u = nn.Parameter(torch.empty(lora_rank, out_features))  # transposed
-        self.v = nn.Parameter(torch.empty(in_features, lora_rank))
-        self.sigma = nn.Parameter(torch.empty(lora_rank))
+        # Meta tensors cannot be copied-from; keep original params and rely on later init.
+        if linear.weight.is_meta:
+            module.weight = linear.weight
+            module.bias = linear.bias
+        else:
+            module.weight.data.copy_(linear.weight.data)
+            if linear.bias is not None:
+                module.bias.data.copy_(linear.bias.data)
+
+        # Initialize LoRA parameters (matches the previous init_lora_weights defaults)
+        if not module.u.is_meta:
+            nn.init.normal_(module.u, mean=0.0, std=0.02)
+            nn.init.zeros_(module.v)
+            nn.init.ones_(module.sigma)
+
+        return module
 
     def forward(self, input):
         lora_update = (input @ self.v) * self.sigma
@@ -25,20 +68,3 @@ def forward(self, input):
         if self.bias is not None:
             y += self.bias
         return y
-
-    @classmethod
-    def from_linear(cls, linear: nn.Linear, lora_rank: int = 32):
-        new_layer = cls(linear.in_features, linear.out_features, linear.bias is not None, lora_rank)
-        new_layer.weight = linear.weight
-        new_layer.bias = linear.bias
-        device = linear.weight.device
-        dtype = linear.weight.dtype
-        new_layer.u.data = new_layer.u.data.to(device=device, dtype=dtype)
-        new_layer.v.data = new_layer.v.data.to(device=device, dtype=dtype)
-        new_layer.sigma.data = new_layer.sigma.data.to(device=device, dtype=dtype)
-        return new_layer
-    
-    def init_lora_weights(self, init_std: float = 0.02):
-        nn.init.normal_(self.u, mean=0.0, std=init_std)
-        nn.init.zeros_(self.v)
-        nn.init.ones_(self.sigma)
diff --git a/alto/observers/per_channel_norm.py b/alto/observers/per_channel_norm.py
index 6a1cda79..7c49c126 100644
--- a/alto/observers/per_channel_norm.py
+++ b/alto/observers/per_channel_norm.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MIT
 
 import torch
+from torch.distributed.tensor import DTensor, Partial
 from alto.utils.pytorch.module import TransformerConv1D
 from .base import Observer, register_observer
 
@@ -25,7 +26,10 @@ def __init__(self, *args, **kwargs) -> None:
     def make_empty_row_scalars(self) -> torch.Tensor:
         weight = self.module().weight
         num_columns = weight.shape[1]
-        return torch.zeros(num_columns, device=self.device)
+        S = torch.zeros(num_columns, device=self.device)
+        if isinstance(weight, DTensor):
+            S = DTensor.from_local(S, device_mesh=weight.device_mesh, placements=(Partial("avg"),))
+        return S
 
     def get_current_min_max(self, observed: torch.Tensor):
         pass
@@ -38,6 +42,13 @@ def forward_inner(self, x_orig):
             return x_orig
 
         with torch.no_grad():
+            S = self.stats
+            if isinstance(S, DTensor):
+                S = S.to_local()
+
+            # TODO: support TP
+            assert not isinstance(x_orig, DTensor), "TP is not supported for per_channel_norm observer"
+
             module = self.module()
             inp = x_orig.detach()
             if inp.dim() == 2:
@@ -62,11 +73,11 @@ def forward_inner(self, x_orig):
                 inp = inp.permute([1, 0, 2])
                 inp = inp.flatten(1)
 
-            self.stats *= self.num_samples / (self.num_samples + num_added)
+            S *= self.num_samples / (self.num_samples + num_added)
             self.num_samples += num_added
 
             inp = inp.type(PRECISION)
-            self.stats += torch.norm(inp, p=2, dim=1)**2 / self.num_samples
+            S += torch.norm(inp, p=2, dim=1)**2 / self.num_samples
 
         return x_orig
 
diff --git a/alto/train.py b/alto/train.py
index cec58aec..0b12a350 100644
--- a/alto/train.py
+++ b/alto/train.py
@@ -148,9 +148,10 @@ def forward_step(
         model_parts = self.model_parts
         parallel_dims = self.parallel_dims
 
-        inputs, _, extra_inputs, extra_kwargs = self.post_dataloading_process(input_dict, labels)
+        inputs, labels, extra_kwargs = self.post_dataloading_process(input_dict, labels)
 
         if parallel_dims.pp_enabled:
+            loss_kwargs = {"global_valid_tokens": global_valid_tokens}
             targets, losses = None, None
             result = None
             with self.train_context():
@@ -158,16 +159,18 @@ def forward_step(
                     if self.pp_has_first_stage:
                         self.pp_schedule.eval(
                             inputs,
-                            **extra_inputs,
                             **extra_kwargs,
                             target=targets,
                             losses=losses,
+                            loss_kwargs=loss_kwargs,
+                            return_outputs=False,
                         )
                     elif self.pp_has_last_stage:
                         result = self.pp_schedule.eval(
                             **extra_kwargs,
                             target=targets,
                             losses=losses,
+                            loss_kwargs=loss_kwargs,
                             return_outputs=True,
                         )
                     else:
@@ -175,13 +178,14 @@ def forward_step(
                             **extra_kwargs,
                             target=targets,
                             losses=losses,
+                            loss_kwargs=loss_kwargs,
+                            return_outputs=False,
                         )
         else:
             # Non-PP forward / backward
             with self.train_context():
                 assert len(model_parts) == 1
-                with self.maybe_enable_amp:
-                    result = model_parts[0](inputs, **extra_inputs, **extra_kwargs)
+                result = model_parts[0](inputs, **extra_kwargs)
 
         return result
 
diff --git a/alto/utils/exportation/export.py b/alto/utils/exportation/export.py
index e728e47a..e803e4b9 100644
--- a/alto/utils/exportation/export.py
+++ b/alto/utils/exportation/export.py
@@ -60,17 +60,6 @@ def patched_finfo(dtype):
     torch.finfo = orig_finfo_func
 
 
-def hot_fix_for_tied_word_embeddings(model: torch.nn.Module, compressor: ModelCompressor):
-    if not getattr(model.config, "enable_weight_tying", False):
-        return
-    # in the current impl, lm_head is not a linear layer,
-    # so we need to add it to the ignore list manually
-    if compressor.quantization_config is not None:
-        compressor.quantization_config.ignore.append("lm_head")
-    if compressor.sparsity_config is not None:
-        compressor.sparsity_config.ignore.append("lm_head")
-
-
 def get_model_compressor(
     model: torch.nn.Module,
     state_dict: dict[str, torch.Tensor],
@@ -205,7 +194,6 @@ def convert_to_hf(
                     compressor.quantization_config.ignore)
             if compressor.sparsity_config is not None:
                 compressor.sparsity_config.ignore = sd_adapter.map_ignore_list_to_hf(compressor.sparsity_config.ignore)
-            hot_fix_for_tied_word_embeddings(model, compressor)
 
             state_dict = compressor.compress(model)
             compressor.update_config(output_dir)
diff --git a/examples/ADMM_Llama_3.1_8B.sh b/examples/ADMM_Llama_3.1_8B.sh
index ea64fee4..61c6f992 100755
--- a/examples/ADMM_Llama_3.1_8B.sh
+++ b/examples/ADMM_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/ALPS_Llama_3.1_8B.sh b/examples/ALPS_Llama_3.1_8B.sh
index 7e839754..208f83df 100755
--- a/examples/ALPS_Llama_3.1_8B.sh
+++ b/examples/ALPS_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/AWQ_Llama_3.1_8B.sh b/examples/AWQ_Llama_3.1_8B.sh
index dde18ae4..4fcf6ebd 100755
--- a/examples/AWQ_Llama_3.1_8B.sh
+++ b/examples/AWQ_Llama_3.1_8B.sh
@@ -29,6 +29,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/AdmmStructured_Llama_3.1_8B.sh b/examples/AdmmStructured_Llama_3.1_8B.sh
index 472d28d9..54ab2992 100755
--- a/examples/AdmmStructured_Llama_3.1_8B.sh
+++ b/examples/AdmmStructured_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/CosineSimilarity_Llama_3.1_8B.sh b/examples/CosineSimilarity_Llama_3.1_8B.sh
index 0dc6e180..335f4a4d 100755
--- a/examples/CosineSimilarity_Llama_3.1_8B.sh
+++ b/examples/CosineSimilarity_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/GPTQ_Llama_3.1_8B.sh b/examples/GPTQ_Llama_3.1_8B.sh
index 1296fd9f..22733aa4 100755
--- a/examples/GPTQ_Llama_3.1_8B.sh
+++ b/examples/GPTQ_Llama_3.1_8B.sh
@@ -28,6 +28,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/Magnitude_Llama_3.1_8B.sh b/examples/Magnitude_Llama_3.1_8B.sh
index d6e37870..b00bcae4 100755
--- a/examples/Magnitude_Llama_3.1_8B.sh
+++ b/examples/Magnitude_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/OBS_Llama_3.1_8B.sh b/examples/OBS_Llama_3.1_8B.sh
index 29bc5bbd..93c64eee 100755
--- a/examples/OBS_Llama_3.1_8B.sh
+++ b/examples/OBS_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/RTN_Llama_3.1_8B.sh b/examples/RTN_Llama_3.1_8B.sh
index a9be4df7..7dbe1c5b 100755
--- a/examples/RTN_Llama_3.1_8B.sh
+++ b/examples/RTN_Llama_3.1_8B.sh
@@ -41,6 +41,7 @@ if [ -n "$COMM_MODE" ]; then
 else
     # Normal training with torchrun
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/SparseGPT_Llama_3.1_8B.sh b/examples/SparseGPT_Llama_3.1_8B.sh
index 075c48db..2f708cb5 100755
--- a/examples/SparseGPT_Llama_3.1_8B.sh
+++ b/examples/SparseGPT_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/WandaStructured_Llama_3.1_8B.sh b/examples/WandaStructured_Llama_3.1_8B.sh
index 385740fb..4d47fae3 100755
--- a/examples/WandaStructured_Llama_3.1_8B.sh
+++ b/examples/WandaStructured_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/Wanda_Llama_3.1_8B.sh b/examples/Wanda_Llama_3.1_8B.sh
index 4547a654..c7fa9301 100755
--- a/examples/Wanda_Llama_3.1_8B.sh
+++ b/examples/Wanda_Llama_3.1_8B.sh
@@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then
     NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/examples/llama3.2_1b_mx9.sh b/examples/llama3.2_1b_mx9.sh
index 204794a6..59bb221f 100755
--- a/examples/llama3.2_1b_mx9.sh
+++ b/examples/llama3.2_1b_mx9.sh
@@ -49,6 +49,7 @@ if [ -n "$COMM_MODE" ]; then
 else
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
         --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
         -m ${TRAIN_FILE} \
diff --git a/examples/run.sh b/examples/run.sh
index 4c6ff0c6..2957492a 100755
--- a/examples/run.sh
+++ b/examples/run.sh
@@ -44,6 +44,7 @@ if [ -n "$COMM_MODE" ]; then
 else
     # Normal training with torchrun
     PYTORCH_ALLOC_CONF="expandable_segments:True" \
+    HSA_NO_SCRATCH_RECLAIM=1 \
     TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \
     torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
diff --git a/tests/integration/gpt_oss_debugmodel_lpt.sh b/tests/integration/gpt_oss_debugmodel_lpt.sh
index 7037a540..fb568804 100755
--- a/tests/integration/gpt_oss_debugmodel_lpt.sh
+++ b/tests/integration/gpt_oss_debugmodel_lpt.sh
@@ -6,4 +6,5 @@ cd $SCRIPT_DIR/../..
 NGPU=2 \
 MODULE=gpt_oss \
 CONFIG=gpt_oss_debugmodel_lpt \
-./examples/run.sh
+./examples/run.sh \
+  --parallelism.expert_parallel_degree 2
diff --git a/tests/integration/instella_3b_opt.sh b/tests/integration/instella_3b_opt.sh
deleted file mode 100755
index c7d4ed13..00000000
--- a/tests/integration/instella_3b_opt.sh
+++ /dev/null
@@ -1,11 +0,0 @@
-#!/bin/bash
-
-SCRIPT_DIR=$(dirname "$0")
-cd $SCRIPT_DIR/../..
-
-NGPU=1 \
-MODULE=llama3 \
-CONFIG=instella_3b_opt \
-./examples/run.sh \
-  --hf_assets_path=/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e \
-  --checkpoint.initial_load_path=/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e
diff --git a/tests/unittest/nn/test_decomposed_linear.py b/tests/unittest/nn/test_decomposed_linear.py
index 99376d85..191540fa 100644
--- a/tests/unittest/nn/test_decomposed_linear.py
+++ b/tests/unittest/nn/test_decomposed_linear.py
@@ -41,7 +41,7 @@ def test_decomposed_linear(in_features, out_features, bias, lora_rank):
 @pytest.mark.parametrize("precision", ["mxfp4"])
 def test_decomposed_linear_quantization(in_features, out_features, bias, lora_rank, precision):
     STD = 0.1
-    decomposed_linear = DecomposedLinear(in_features, out_features, bias, lora_rank).to("cuda")
+    decomposed_linear = DecomposedLinear(DecomposedLinear.Config(in_features=in_features, out_features=out_features, bias=bias, lora_rank=lora_rank)).to("cuda")
     decomposed_linear.weight.data.normal_(mean=0, std=STD)
     if bias:
         decomposed_linear.bias.data.normal_(mean=0, std=STD)
diff --git a/version.txt b/version.txt
index 22701ea2..206c0852 100644
--- a/version.txt
+++ b/version.txt
@@ -1 +1 @@
-0.0.2-dev0
+0.1.0-dev0