From 64f8231a3ed8e762fd58a75bcb808e17b96a6317 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 26 May 2026 10:31:39 -0700 Subject: [PATCH 1/7] Fix TEGroupedMLP pre-backward unshard in fine-grained FSDP hooks for 1F1B EP overlap. Signed-off-by: Ritesh Patel --- .../distributed/fsdp/mcore_fsdp_adapter.py | 22 ++++++++++++++++++- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 22 ++++++++++++++++++- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 6c7ec1c5bd7..e53eb337259 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -14,7 +14,7 @@ import logging import random -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple, Type try: import einops @@ -41,6 +41,7 @@ from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_layer import MambaLayer +from megatron.core.transformer.moe.experts import TEGroupedMLP from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import MoETransformerLayer, TransformerLayer from megatron.core.utils import is_te_min_version, log_single_rank @@ -91,6 +92,20 @@ class FullyShardedDataParallel(_BaseDataParallel): }, } + @staticmethod + def _fine_grained_pre_backward_recurse_module_types( + config: TransformerConfig, + ddp_config: DistributedDataParallelConfig, + ) -> Tuple[Type[nn.Module], ...]: + """Module classes that need ``parameters(recurse=True)`` for overlap hooks.""" + if ( + config.overlap_moe_expert_parallel_comm + and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" + ): + # Subclasses (e.g. InferenceGroupedMLP) are covered by isinstance. + return (TEGroupedMLP,) + return () + def __init__( self, config: TransformerConfig, @@ -211,6 +226,11 @@ def __init__( config.overlap_moe_expert_parallel_comm and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" ), + fine_grained_pre_backward_recurse_module_types=( + self._fine_grained_pre_backward_recurse_module_types( + config, ddp_config + ) + ), ), ) self.param_and_grad_buffer = self.module.param_and_grad_buffer diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index e90e3879dcb..5ca60f8350b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from enum import Enum, auto from functools import partial -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Type import torch import torch.nn as nn @@ -177,6 +177,12 @@ class MegatronFSDP(torch.nn.Module): userbuffer registration when nccl_ub is set. enable_fine_grained_param_gather (bool): Whether to enable "fine-grained" param all-gather, which can improve performance when using MXFP8 parameters with activation recomputation. + enable_fine_grained_param_gather_backward_hook (bool): Register pre-backward unshard hooks + on each submodule (used by 1F1B EP overlap and similar schedules). + fine_grained_pre_backward_recurse_module_types (Optional[Tuple[Type[nn.Module], ...]]): + Module classes for which fine-grained pre-backward unshard uses + ``parameters(recurse=True)`` (container modules whose sharded weights live on + children). Checked with :func:`isinstance`. Defaults to empty (none). report_nan_in_param_grad (bool): Whether to enable precise NaN-checking for parameter wgrad. Can significantly degrade performance. Defaults to False. @@ -217,6 +223,9 @@ def __init__( disable_symmetric_registration: bool = False, enable_fine_grained_param_gather_hook: bool = False, enable_fine_grained_param_gather_backward_hook: bool = False, + fine_grained_pre_backward_recurse_module_types: Optional[ + Tuple[Type[nn.Module], ...] + ] = None, report_nan_in_param_grad: bool = False, ): super().__init__() @@ -272,6 +281,10 @@ def __init__( self.enable_fine_grained_param_gather_backward_hook = ( enable_fine_grained_param_gather_backward_hook ) + recurse_types = fine_grained_pre_backward_recurse_module_types or () + self.fine_grained_pre_backward_recurse_module_types: Tuple[ + Type[nn.Module], ... + ] = recurse_types self.report_nan_in_param_grad = report_nan_in_param_grad # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP. @@ -863,6 +876,13 @@ def _pre_backward_param_unshard(module: nn.Module, *unused): if isinstance(module, tuple(fsdp_unit_modules)): param_list = list(module.parameters()) + elif ( + self.enable_fine_grained_param_gather_backward_hook + and self.fine_grained_pre_backward_recurse_module_types + and isinstance(module, self.fine_grained_pre_backward_recurse_module_types) + ): + # Container modules (e.g. TEGroupedMLP): sharded weights are on children. + param_list = list(module.parameters(recurse=True)) else: param_list = list(module.parameters(recurse=False)) From 2819a1792a5a8aaa3f92da6984a536a7143b6225 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Thu, 28 May 2026 13:57:14 -0700 Subject: [PATCH 2/7] Symmetric fine-grained FSDP unshard for TEGroupedMLP via fine_grained_recurse_module_types Signed-off-by: Ritesh Patel --- .../distributed/fsdp/mcore_fsdp_adapter.py | 10 ++-- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 56 +++++++++---------- 2 files changed, 30 insertions(+), 36 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index e53eb337259..9629566cd08 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -93,11 +93,11 @@ class FullyShardedDataParallel(_BaseDataParallel): } @staticmethod - def _fine_grained_pre_backward_recurse_module_types( + def _fine_grained_recurse_module_types( config: TransformerConfig, ddp_config: DistributedDataParallelConfig, ) -> Tuple[Type[nn.Module], ...]: - """Module classes that need ``parameters(recurse=True)`` for overlap hooks.""" + """Container module classes that need ``parameters(recurse=True)`` for fine-grained hooks.""" if ( config.overlap_moe_expert_parallel_comm and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" @@ -226,10 +226,8 @@ def __init__( config.overlap_moe_expert_parallel_comm and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" ), - fine_grained_pre_backward_recurse_module_types=( - self._fine_grained_pre_backward_recurse_module_types( - config, ddp_config - ) + fine_grained_recurse_module_types=self._fine_grained_recurse_module_types( + config, ddp_config ), ), ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 5ca60f8350b..33374d4c4f2 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -179,8 +179,8 @@ class MegatronFSDP(torch.nn.Module): which can improve performance when using MXFP8 parameters with activation recomputation. enable_fine_grained_param_gather_backward_hook (bool): Register pre-backward unshard hooks on each submodule (used by 1F1B EP overlap and similar schedules). - fine_grained_pre_backward_recurse_module_types (Optional[Tuple[Type[nn.Module], ...]]): - Module classes for which fine-grained pre-backward unshard uses + fine_grained_recurse_module_types (Optional[Tuple[Type[nn.Module], ...]]): + Module classes for which fine-grained pre-forward / pre-backward unshard uses ``parameters(recurse=True)`` (container modules whose sharded weights live on children). Checked with :func:`isinstance`. Defaults to empty (none). report_nan_in_param_grad (bool): Whether to enable precise NaN-checking for parameter wgrad. @@ -223,9 +223,7 @@ def __init__( disable_symmetric_registration: bool = False, enable_fine_grained_param_gather_hook: bool = False, enable_fine_grained_param_gather_backward_hook: bool = False, - fine_grained_pre_backward_recurse_module_types: Optional[ - Tuple[Type[nn.Module], ...] - ] = None, + fine_grained_recurse_module_types: Optional[Tuple[Type[nn.Module], ...]] = None, report_nan_in_param_grad: bool = False, ): super().__init__() @@ -281,10 +279,8 @@ def __init__( self.enable_fine_grained_param_gather_backward_hook = ( enable_fine_grained_param_gather_backward_hook ) - recurse_types = fine_grained_pre_backward_recurse_module_types or () - self.fine_grained_pre_backward_recurse_module_types: Tuple[ - Type[nn.Module], ... - ] = recurse_types + recurse_types = fine_grained_recurse_module_types or () + self.fine_grained_recurse_module_types: Tuple[Type[nn.Module], ...] = recurse_types self.report_nan_in_param_grad = report_nan_in_param_grad # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP. @@ -559,6 +555,25 @@ def _register_fsdp_hooks(self, root_module): """ fsdp_unit_modules = self.fsdp_unit_modules + def _param_list_for_submodule_unshard( + module: nn.Module, *, for_backward: bool + ) -> List[nn.Parameter]: + """Build the parameter list for fine-grained or FSDP-unit unshard hooks.""" + if isinstance(module, tuple(fsdp_unit_modules)): + return list(module.parameters()) + fine_grained_enabled = ( + self.enable_fine_grained_param_gather_backward_hook + if for_backward + else self.enable_fine_grained_param_gather_hook + ) + if ( + fine_grained_enabled + and self.fine_grained_recurse_module_types + and isinstance(module, self.fine_grained_recurse_module_types) + ): + return list(module.parameters(recurse=True)) + return list(module.parameters(recurse=False)) + def release_module_parameters(module, bwd, lazy=False, *unused): """ Release the parameters of a given module after completing the forward @@ -749,16 +764,7 @@ def _pre_forward_param_unshard(module: nn.Module, *unused): else: module._training_state = TrainingState.FORWARD - if isinstance(module, tuple(fsdp_unit_modules)): - param_list = list(module.parameters()) - else: - # All-gather the shallow parameters in every forward pass for modules - # that are not FSDP units. Do not recurse unless absolutely necessary, - # to allocate as little memory as possible for this forward pass. - param_list = list(module.parameters(recurse=False)) - - if self.enable_fine_grained_param_gather_hook: - param_list = list(module.parameters(recurse=False)) + param_list = _param_list_for_submodule_unshard(module, for_backward=False) # All-gather the parameters before the forward pass. self.all_gather_and_wait_parameters_ready( @@ -874,17 +880,7 @@ def _pre_backward_param_unshard(module: nn.Module, *unused): for sub_module in module.modules(): sub_module._training_state = TrainingState.PRE_BACKWARD - if isinstance(module, tuple(fsdp_unit_modules)): - param_list = list(module.parameters()) - elif ( - self.enable_fine_grained_param_gather_backward_hook - and self.fine_grained_pre_backward_recurse_module_types - and isinstance(module, self.fine_grained_pre_backward_recurse_module_types) - ): - # Container modules (e.g. TEGroupedMLP): sharded weights are on children. - param_list = list(module.parameters(recurse=True)) - else: - param_list = list(module.parameters(recurse=False)) + param_list = _param_list_for_submodule_unshard(module, for_backward=True) # All-gather / unshard the module parameters before the backward pass. self.all_gather_and_wait_parameters_ready( From 0ad5777244f6dddeba6d90bf3aad616ea5fbcb06 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Thu, 28 May 2026 14:53:50 -0700 Subject: [PATCH 3/7] Rename submodule unshard helper arg to pass_direction Literal[forward, backward]. Signed-off-by: Ritesh Patel --- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 33374d4c4f2..a917e8e9f72 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from enum import Enum, auto from functools import partial -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import Any, Dict, List, Literal, Optional, Tuple, Type import torch import torch.nn as nn @@ -556,14 +556,14 @@ def _register_fsdp_hooks(self, root_module): fsdp_unit_modules = self.fsdp_unit_modules def _param_list_for_submodule_unshard( - module: nn.Module, *, for_backward: bool + module: nn.Module, pass_direction: Literal["forward", "backward"] ) -> List[nn.Parameter]: """Build the parameter list for fine-grained or FSDP-unit unshard hooks.""" if isinstance(module, tuple(fsdp_unit_modules)): return list(module.parameters()) fine_grained_enabled = ( self.enable_fine_grained_param_gather_backward_hook - if for_backward + if pass_direction == "backward" else self.enable_fine_grained_param_gather_hook ) if ( @@ -764,7 +764,7 @@ def _pre_forward_param_unshard(module: nn.Module, *unused): else: module._training_state = TrainingState.FORWARD - param_list = _param_list_for_submodule_unshard(module, for_backward=False) + param_list = _param_list_for_submodule_unshard(module, "forward") # All-gather the parameters before the forward pass. self.all_gather_and_wait_parameters_ready( @@ -880,7 +880,7 @@ def _pre_backward_param_unshard(module: nn.Module, *unused): for sub_module in module.modules(): sub_module._training_state = TrainingState.PRE_BACKWARD - param_list = _param_list_for_submodule_unshard(module, for_backward=True) + param_list = _param_list_for_submodule_unshard(module, "backward") # All-gather / unshard the module parameters before the backward pass. self.all_gather_and_wait_parameters_ready( From 50672eb3c5da29652423aa917a1c592b3052e3da Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Wed, 24 Jun 2026 16:21:21 -0700 Subject: [PATCH 4/7] Lazy-import TEGroupedMLP to fix fsdp_dtensor checkpoint load. Signed-off-by: Ritesh Patel --- megatron/core/distributed/fsdp/mcore_fsdp_adapter.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 9629566cd08..49ae276cf69 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -41,7 +41,6 @@ from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_layer import MambaLayer -from megatron.core.transformer.moe.experts import TEGroupedMLP from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import MoETransformerLayer, TransformerLayer from megatron.core.utils import is_te_min_version, log_single_rank @@ -102,7 +101,9 @@ def _fine_grained_recurse_module_types( config.overlap_moe_expert_parallel_comm and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" ): - # Subclasses (e.g. InferenceGroupedMLP) are covered by isinstance. + # Lazy import to avoid circular chain. + from megatron.core.transformer.moe.experts import TEGroupedMLP + return (TEGroupedMLP,) return () From c0422cbdef5e5b1933af10fcd9d19d171d794e86 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Fri, 26 Jun 2026 15:56:30 -0700 Subject: [PATCH 5/7] Fix fine-grained FSDP unshard logic and add SharedExpertMLP to recurse allowlist. Signed-off-by: Ritesh Patel --- .../distributed/fsdp/mcore_fsdp_adapter.py | 3 +- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 44 ++++++++++++++----- 2 files changed, 36 insertions(+), 11 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 49ae276cf69..276c40b8c5e 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -103,8 +103,9 @@ def _fine_grained_recurse_module_types( ): # Lazy import to avoid circular chain. from megatron.core.transformer.moe.experts import TEGroupedMLP + from megatron.core.transformer.moe.shared_experts import SharedExpertMLP - return (TEGroupedMLP,) + return (TEGroupedMLP, SharedExpertMLP) return () def __init__( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index a917e8e9f72..51b593464fe 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -558,21 +558,45 @@ def _register_fsdp_hooks(self, root_module): def _param_list_for_submodule_unshard( module: nn.Module, pass_direction: Literal["forward", "backward"] ) -> List[nn.Parameter]: - """Build the parameter list for fine-grained or FSDP-unit unshard hooks.""" - if isinstance(module, tuple(fsdp_unit_modules)): - return list(module.parameters()) + """Build the parameter list for fine-grained or FSDP-unit unshard hooks. + + Parameter buckets designated by this function are all-gathered and may + pre-fetch subsequent buckets in FSDP bucket order during runtime. + """ + # Fine-grained hooks are attached to all sub-modules; this function + # controls which parameters each hook should unshard. fine_grained_enabled = ( self.enable_fine_grained_param_gather_backward_hook if pass_direction == "backward" else self.enable_fine_grained_param_gather_hook ) - if ( - fine_grained_enabled - and self.fine_grained_recurse_module_types - and isinstance(module, self.fine_grained_recurse_module_types) - ): - return list(module.parameters(recurse=True)) - return list(module.parameters(recurse=False)) + if fine_grained_enabled: + # Fine-grained hooks run on every submodule: shallow params by + # default, including on FSDP units (e.g. TransformerLayer). Leaf + # child hooks gather their own nested weights. Container modules + # in fine_grained_recurse_module_types (e.g. TEGroupedMLP, + # SharedExpertMLP) need recurse=True because weights live on + # children and the container is the compute entry point. + if ( + self.fine_grained_recurse_module_types + and isinstance(module, self.fine_grained_recurse_module_types) + ): + return list(module.parameters(recurse=True)) + else: + # Only unshard direct parameters. Used when submodules are + # called in isolation of an FSDP-unit forward (e.g. mxfp8 + # param gather, EP-overlap 1F1B schedule). Leaf modules + # (e.g. TELinear) still gather their own weights via + # separate hooks. Also limits unshard scope for activation + # recomputation on individual submodules. + return list(module.parameters(recurse=False)) + else: + if isinstance(module, tuple(fsdp_unit_modules)): + # FSDP unit modules should be unsharded and communicated together. + return list(module.parameters()) + else: + # Non-unit modules should only unshard the direct parameters they need. + return list(module.parameters(recurse=False)) def release_module_parameters(module, bwd, lazy=False, *unused): """ From 00ddbc67074c40ba037bf711c6e8ec03a976ad21 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 30 Jun 2026 11:37:55 -0700 Subject: [PATCH 6/7] Applying formatting. Signed-off-by: Ritesh Patel --- megatron/core/distributed/fsdp/mcore_fsdp_adapter.py | 3 +-- .../core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py | 5 ++--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 276c40b8c5e..51e4ff0956b 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -93,8 +93,7 @@ class FullyShardedDataParallel(_BaseDataParallel): @staticmethod def _fine_grained_recurse_module_types( - config: TransformerConfig, - ddp_config: DistributedDataParallelConfig, + config: TransformerConfig, ddp_config: DistributedDataParallelConfig ) -> Tuple[Type[nn.Module], ...]: """Container module classes that need ``parameters(recurse=True)`` for fine-grained hooks.""" if ( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 51b593464fe..58f4a2d8206 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -577,9 +577,8 @@ def _param_list_for_submodule_unshard( # in fine_grained_recurse_module_types (e.g. TEGroupedMLP, # SharedExpertMLP) need recurse=True because weights live on # children and the container is the compute entry point. - if ( - self.fine_grained_recurse_module_types - and isinstance(module, self.fine_grained_recurse_module_types) + if self.fine_grained_recurse_module_types and isinstance( + module, self.fine_grained_recurse_module_types ): return list(module.parameters(recurse=True)) else: From cc78b379a420e37fcb68c35d08723c54ea7295a9 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 30 Jun 2026 14:47:00 -0700 Subject: [PATCH 7/7] Fix line-too-long linting error. Signed-off-by: Ritesh Patel --- megatron/core/distributed/fsdp/mcore_fsdp_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 51e4ff0956b..efabf4e58ee 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -95,7 +95,7 @@ class FullyShardedDataParallel(_BaseDataParallel): def _fine_grained_recurse_module_types( config: TransformerConfig, ddp_config: DistributedDataParallelConfig ) -> Tuple[Type[nn.Module], ...]: - """Container module classes that need ``parameters(recurse=True)`` for fine-grained hooks.""" + """Module classes needing ``parameters(recurse=True)`` for fine-grained hooks.""" if ( config.overlap_moe_expert_parallel_comm and ddp_config.data_parallel_sharding_strategy == "optim_grads_params"