Specify Trainer strategy in CLI config?

[wwbrannon]

Hi folks,

I can't figure out how to give a Strategy class (rather than a string) as an argument to the trainer in a CLI config file. Is doing so supported?

Doing this works fine:

trainer:
    ...
    strategy: 'deepspeed_stage_2_offload'
    ...

But this gives a jsonargparse error:

trainer:
    ...
    strategy:
        class_path: pytorch_lightning.strategies.DeepSpeedStrategy
        init_args:
            stage: 2
            offload_optimizer: True
            logging_batch_size_per_gpu: 16
    ...

The error:

Traceback (most recent call last):
  File "/media/sharon/wbrannon/github/clip-graph/bin/trainer.py", line 7, in <module>
    cli = cl.LightningCLI(
  File "/media/sharon/wbrannon/miniconda3/envs/dl/lib/python3.9/site-packages/pytorch_lightning/utilities/cli.py", line 552, in __init__
    self.parse_arguments(self.parser)
  File "/media/sharon/wbrannon/miniconda3/envs/dl/lib/python3.9/site-packages/pytorch_lightning/utilities/cli.py", line 692, in parse_arguments
    self.config = parser.parse_args()
  File "/media/sharon/wbrannon/miniconda3/envs/dl/lib/python3.9/site-packages/pytorch_lightning/utilities/cli.py", line 268, in parse_args
    return super().parse_args(*args, **kwargs)
  File "/media/sharon/wbrannon/miniconda3/envs/dl/lib/python3.9/site-packages/jsonargparse/deprecated.py", line 112, in patched_parse
    cfg = parse_method(*args, _skip_check=_skip_check, **kwargs)
  File "/media/sharon/wbrannon/miniconda3/envs/dl/lib/python3.9/site-packages/jsonargparse/core.py", line 366, in parse_args
    self.error(str(ex), ex)
  File "/media/sharon/wbrannon/miniconda3/envs/dl/lib/python3.9/site-packages/jsonargparse/core.py", line 983, in error
    raise ParserError(message) from ex
jsonargparse.util.ParserError: Parser key "trainer.strategy": Value "Namespace(class_path='pytorch_lightning.strategies.DeepSpeedStrategy', init_args=Namespace(zero_optimization=True, s
tage=2, remote_device='cpu', offload_optimizer=False, offload_parameters=False, offload_params_device='cpu', nvme_path='/local_nvme', params_buffer_count=5, params_buffer_size=100000000
.0, max_in_cpu=1000000000.0, offload_optimizer_device='cpu', optimizer_buffer_count=4, block_size=1048576, queue_depth=8, single_submit=False, overlap_events=True, thread_count=1, pin_memory=False, sub_group_size=1000000000000.0, contiguous_gradients=True, overlap_comm=True, allgather_partitions=True, reduce_scatter=True, allgather_bucket_size=200000000.0, reduce_bucket_size=200000000.0, zero_allow_untested_optimizer=True, logging_batch_size_per_gpu='auto', config=None, logging_level=30, parallel_devices=None, cluster_environment=None, loss_scale=0.0, initial_scale_power=16, loss_scale_window=1000, hysteresis=2, min_loss_scale=1, partition_activations=False, cpu_checkpointing=False, contiguous_memory_optimization=False, synchronize_checkpoint_boundary=False, load_full_weights=False, precision_plugin=None, process_group_backend=None))" does not validate against any of the types in typing.Union[str, pytorch_lightning.strategies.strategy.Strategy, NoneType]:
  - Expected a <class 'str'> but got "Namespace(class_path='pytorch_lightning.strategies.DeepSpeedStrategy', init_args=Namespace(zero_optimization=True, stage=2, remote_device='cpu', offload_optimizer=False, offload_parameters=False, offload_params_device='cpu', nvme_path='/local_nvme', params_buffer_count=5, params_buffer_size=100000000.0, max_in_cpu=1000000000.0, offload_optimizer_device='cpu', optimizer_buffer_count=4, block_size=1048576, queue_depth=8, single_submit=False, overlap_events=True, thread_count=1, pin_memory=False, sub_group_size=1000000000000.0, contiguous_gradients=True, overlap_comm=True, allgather_partitions=True, reduce_scatter=True, allgather_bucket_size=200000000.0, reduce_bucket_size=200000000.0, zero_allow_untested_optimizer=True, logging_batch_size_per_gpu='auto', config=None, logging_level=30, parallel_devices=None, cluster_environment=None, loss_scale=0.0, initial_scale_power=16, loss_scale_window=1000, hysteresis=2, min_loss_scale=1, partition_activations=False, cpu_checkpointing=False, contiguous_memory_optimization=False, synchronize_checkpoint_boundary=False, load_full_weights=False, precision_plugin=None, process_group_backend=None))"
  - Problem with given class_path "pytorch_lightning.strategies.DeepSpeedStrategy":
    - Configuration check failed :: Parser key "params_buffer_size": Expected a <class 'int'> but got "100000000.0"
  - Expected a <class 'NoneType'> but got "Namespace(class_path='pytorch_lightning.strategies.DeepSpeedStrategy', init_args=Namespace(zero_optimization=True, stage=2, remote_device='cpu', offload_optimizer=False, offload_parameters=False, offload_params_device='cpu', nvme_path='/local_nvme', params_buffer_count=5, params_buffer_size=100000000.0, max_in_cpu=1000000000.0, offload_optimizer_device='cpu', optimizer_buffer_count=4, block_size=1048576, queue_depth=8, single_submit=False, overlap_events=True, thread_count=1, pin_memory=False, sub_group_size=1000000000000.0, contiguous_gradients=True, overlap_comm=True, allgather_partitions=True, reduce_scatter=True, allgather_bucket_size=200000000.0, reduce_bucket_size=200000000.0, zero_allow_untested_optimizer=True, logging_batch_size_per_gpu='auto', config=None, logging_level=30, parallel_devices=None, cluster_environment=None, loss_scale=0.0, initial_scale_power=16, loss_scale_window=1000, hysteresis=2, min_loss_scale=1, partition_activations=False, cpu_checkpointing=False, contiguous_memory_optimization=False, synchronize_checkpoint_boundary=False, load_full_weights=False, precision_plugin=None, process_group_backend=None))"

[rohitgr7]

cc @carmocca

[carmocca]

It's a bug! Will be fixed with #12989 which should be included in next week's bugfix release

Thanks for the report!