Make enable_hlfb true by default

protobird-git · copybara-github · commit 6a0361efec5e · 2025-05-16T10:25:36.000-07:00
- It's been true for most cases except stable_diffusion's decoder and diffuser
- which use unet config instead of model config and set enable_hlfb explicitly based on cpu/gpu
- Set enable_hlfb false explicitly for cpu_only/gemma3, T5, toy_models, and stable_diffusion clip

PiperOrigin-RevId: 759665348
diff --git a/ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py b/ai_edge_torch/generative/examples/amd_llama_135m/amd_llama_135m.py
@@ -51,9 +51,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=2048,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -69,7 +67,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/deepseek/deepseek.py b/ai_edge_torch/generative/examples/deepseek/deepseek.py
@@ -53,9 +53,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=8960,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-06,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -72,7 +70,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/gemma/gemma1.py b/ai_edge_torch/generative/examples/gemma/gemma1.py
@@ -65,10 +65,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=16384,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -87,7 +84,6 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/gemma/gemma2.py b/ai_edge_torch/generative/examples/gemma/gemma2.py
@@ -233,10 +233,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for a Gemma 2B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -284,7 +281,6 @@ def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
       final_logit_softcap=30.0,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/gemma3/decoder.py b/ai_edge_torch/generative/examples/gemma3/decoder.py
@@ -329,10 +329,7 @@ def get_decoder_config_1b(kv_cache_max_len: int = 2048) -> cfg.ModelConfig:
     The model config for a Gemma 1B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True,
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -379,7 +376,6 @@ def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
       final_logit_softcap=None,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/gemma3/gemma3.py b/ai_edge_torch/generative/examples/gemma3/gemma3.py
@@ -158,9 +158,7 @@ def get_fake_model_config(**kwargs) -> Gemma3MMConfig:
       image_projection_scale=128**0.5,
       image_projection_use_bias=False,
       mm_norm_config=cfg.NormalizationConfig(
-          type=cfg.NormalizationType.LAYER_NORM,
-          epsilon=1e-6,
-          enable_hlfb=True,
+          type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
       ),
       mm_extra_tokens=32,
   )
diff --git a/ai_edge_torch/generative/examples/gemma3/image_encoder.py b/ai_edge_torch/generative/examples/gemma3/image_encoder.py
@@ -98,9 +98,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       output_proj_use_bias=True,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM,
-      epsilon=1e-6,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.SEQUENTIAL,
@@ -123,7 +121,6 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       num_mm_tokens_per_image=256,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/hammer/hammer.py b/ai_edge_torch/generative/examples/hammer/hammer.py
@@ -45,9 +45,7 @@ def get_1_5b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=8960,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-06,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -63,7 +61,6 @@ def get_1_5b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/llama/llama.py b/ai_edge_torch/generative/examples/llama/llama.py
@@ -121,9 +121,7 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True,
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -152,7 +150,6 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/openelm/openelm.py b/ai_edge_torch/generative/examples/openelm/openelm.py
@@ -53,7 +53,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for an OpenELM model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, enable_hlfb=True
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6
   )
   num_heads = [12] * 4 + [16] * 14 + [20] * 12 + [24] * 6
   num_query_groups = [3] * 4 + [4] * 14 + [5] * 12 + [6] * 6
@@ -101,7 +101,6 @@ def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/paligemma/decoder.py b/ai_edge_torch/generative/examples/paligemma/decoder.py
@@ -110,10 +110,7 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=16384,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -132,7 +129,6 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/paligemma/decoder2.py b/ai_edge_torch/generative/examples/paligemma/decoder2.py
@@ -93,10 +93,7 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for the decoder of a PaliGemma 3B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -140,7 +137,6 @@ def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
       final_logit_softcap=30.0,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/paligemma/image_encoder.py b/ai_edge_torch/generative/examples/paligemma/image_encoder.py
@@ -118,9 +118,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       use_bias=True,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM,
-      epsilon=1e-6,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -137,7 +135,6 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/phi/phi2.py b/ai_edge_torch/generative/examples/phi/phi2.py
@@ -66,9 +66,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=10240,
       use_bias=True,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -85,7 +83,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       final_norm_config=norm_config,
       lm_head_use_bias=True,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/phi/phi3.py b/ai_edge_torch/generative/examples/phi/phi3.py
@@ -162,9 +162,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True,
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -192,7 +190,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/phi/phi4.py b/ai_edge_torch/generative/examples/phi/phi4.py
@@ -112,9 +112,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -141,7 +139,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       embedding_dim=3072,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config
diff --git a/ai_edge_torch/generative/examples/qwen/qwen.py b/ai_edge_torch/generative/examples/qwen/qwen.py
@@ -53,9 +53,7 @@ def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=11008,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-06,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -71,7 +69,6 @@ def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/qwen_vl/decoder.py b/ai_edge_torch/generative/examples/qwen_vl/decoder.py
@@ -97,7 +97,7 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=11008,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06, enable_hlfb=True
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -113,7 +113,6 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/qwen_vl/image_encoder.py b/ai_edge_torch/generative/examples/qwen_vl/image_encoder.py
@@ -332,8 +332,7 @@ def get_image_encoder_config(image_size: Tuple[int, int]) -> QwenVLImageConfig:
       use_bias=True,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -359,7 +358,6 @@ def get_image_encoder_config(image_size: Tuple[int, int]) -> QwenVLImageConfig:
       window_size=112,
       spatial_merge_size=2,
       full_atten_block_indexes=[7, 15, 23, 31],
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/smollm/smollm.py b/ai_edge_torch/generative/examples/smollm/smollm.py
@@ -51,9 +51,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=1536,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -68,7 +66,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 
diff --git a/ai_edge_torch/generative/examples/stable_diffusion/clip.py b/ai_edge_torch/generative/examples/stable_diffusion/clip.py
diff --git a/ai_edge_torch/generative/examples/t5/t5.py b/ai_edge_torch/generative/examples/t5/t5.py
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model.py b/ai_edge_torch/generative/examples/test_models/toy_model.py
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py b/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py
diff --git a/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py b/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py
diff --git a/ai_edge_torch/generative/layers/model_config.py b/ai_edge_torch/generative/layers/model_config.py

Original file line number	Diff line number	Diff line change
`@@ -53,9 +53,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:`
`53`	`53`	`intermediate_size=8960,`
`54`	`54`	`)`
`55`	`55`	`norm_config = cfg.NormalizationConfig(`
`56`		`- type=cfg.NormalizationType.RMS_NORM,`
`57`		`- epsilon=1e-06,`
`58`		`- enable_hlfb=True,`
	`56`	`+ type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06`
`59`	`57`	`)`
`60`	`58`	`block_config = cfg.TransformerBlockConfig(`
`61`	`59`	`attn_config=attn_config,`
`@@ -72,7 +70,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:`
`72`	`70`	`block_configs=block_config,`
`73`	`71`	`final_norm_config=norm_config,`
`74`	`72`	`lm_head_share_weight_with_embedding=False,`
`75`		`- enable_hlfb=True,`
`76`	`73`	`)`
`77`	`74`	`return config`
`78`	`75`
Original file line number	Diff line number	Diff line change
`@@ -65,10 +65,7 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:`
`65`	`65`	`intermediate_size=16384,`
`66`	`66`	`)`
`67`	`67`	`norm_config = cfg.NormalizationConfig(`
`68`		`- type=cfg.NormalizationType.RMS_NORM,`
`69`		`- epsilon=1e-6,`
`70`		`- zero_centered=True,`
`71`		`- enable_hlfb=True,`
	`68`	`+ type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True`
`72`	`69`	`)`
`73`	`70`	`block_config = cfg.TransformerBlockConfig(`
`74`	`71`	`attn_config=attn_config,`
`@@ -87,7 +84,6 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:`
`87`	`84`	`block_configs=block_config,`
`88`	`85`	`final_norm_config=norm_config,`
`89`	`86`	`lm_head_use_bias=False,`
`90`		`- enable_hlfb=True,`
`91`	`87`	`)`
`92`	`88`	`return config`
`93`	`89`
Original file line number	Diff line number	Diff line change
`@@ -98,9 +98,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:`
`98`	`98`	`output_proj_use_bias=True,`
`99`	`99`	`)`
`100`	`100`	`norm_config = cfg.NormalizationConfig(`
`101`		`- type=cfg.NormalizationType.LAYER_NORM,`
`102`		`- epsilon=1e-6,`
`103`		`- enable_hlfb=True,`
	`101`	`+ type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6`
`104`	`102`	`)`
`105`	`103`	`ff_config = cfg.FeedForwardConfig(`
`106`	`104`	`type=cfg.FeedForwardType.SEQUENTIAL,`
`@@ -123,7 +121,6 @@ def get_image_encoder_config() -> cfg.ModelConfig:`
`123`	`121`	`image_embedding=image_embedding_config,`
`124`	`122`	`block_configs=block_config,`
`125`	`123`	`final_norm_config=norm_config,`
`126`		`- enable_hlfb=True,`
`127`	`124`	`num_mm_tokens_per_image=256,`
`128`	`125`	`)`
`129`	`126`	`return config`
Original file line number	Diff line number	Diff line change
`@@ -45,9 +45,7 @@ def get_1_5b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:`
`45`	`45`	`intermediate_size=8960,`
`46`	`46`	`)`
`47`	`47`	`norm_config = cfg.NormalizationConfig(`
`48`		`- type=cfg.NormalizationType.RMS_NORM,`
`49`		`- epsilon=1e-06,`
`50`		`- enable_hlfb=True,`
	`48`	`+ type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06`
`51`	`49`	`)`
`52`	`50`	`block_config = cfg.TransformerBlockConfig(`
`53`	`51`	`attn_config=attn_config,`
`@@ -63,7 +61,6 @@ def get_1_5b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:`
`63`	`61`	`kv_cache_max_len=kv_cache_max_len,`
`64`	`62`	`block_configs=block_config,`
`65`	`63`	`final_norm_config=norm_config,`
`66`		`- enable_hlfb=True,`
`67`	`64`	`)`
`68`	`65`	`return config`
`69`	`66`