DDP Sharded with lightning CLI set up does not work: lightning.fabric.utilities.exceptions.MisconfigurationException: Found invalid type for plugin ddp_sharded. Expected one of: Precision, CheckpointIO, ClusterEnviroment, or LayerSync.

[aaronwtr]

I am trying to run my model in a distributed way, but setting the flag plugins: ddp_sharded is failing with the following traceback:

/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/fabric/utilities/seed.py:40: No seed found, seed set to 1982552065
Seed set to 1982552065
Traceback (most recent call last):
  File "/data/home/bty174/OpenBioML/enformer-finetune/main.py", line 16, in <module>
    cli_main()
  File "/data/home/bty174/OpenBioML/enformer-finetune/main.py", line 11, in cli_main
    cli = LightningCLI(EnformerTX, EnformerTXDataModule)
  File "/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/pytorch/cli.py", line 383, in __init__
    self.instantiate_classes()
  File "/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/pytorch/cli.py", line 537, in instantiate_classes
    self.trainer = self.instantiate_trainer()
  File "/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/pytorch/cli.py", line 548, in instantiate_trainer
    return self._instantiate_trainer(trainer_config, extra_callbacks)
  File "/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/pytorch/cli.py", line 573, in _instantiate_trainer
    return self.trainer_class(**config)
  File "/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/pytorch/utilities/argparse.py", line 70, in insert_env_defaults
    return fn(self, **kwargs)
  File "/data/home/bty174/.conda/envs/enformer-finetuning/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 401, in __init__
    self._accelerator_connector = _AcceleratorConnector(

I am using a CLI setup, meaning have a main.py as follows:

def cli_main():
    cli = LightningCLI(EnformerTX, EnformerTXDataModule)


if __name__ == "__main__":
    cli_main()

and a config.yaml where I am setting all the pytorch flags. In the config.yaml I have added the plugins: ddp_sharded flag.

Anyone has any thoughts on what is happening here? I am using lightning version 2.1.1.

[mauvilsa]

I don't know much about the plugins. But from what I can see in the code, the plugins parameter does accept a string as input, which explains why the parser does not complain when given ddp_sharded. However, it doesn't seem to be allowed to give a string, see

pytorch-lightning/src/lightning/pytorch/trainer/connectors/accelerator_connector.py

Lines 245 to 266 in 2f1756f

	if isinstance(plugin, Precision):
	self._precision_plugin_flag = plugin
	plugins_flags_types[Precision.__name__] += 1
	elif isinstance(plugin, CheckpointIO):
	self.checkpoint_io = plugin
	plugins_flags_types[CheckpointIO.__name__] += 1
	elif isinstance(plugin, ClusterEnvironment):
	self._cluster_environment_flag = plugin
	plugins_flags_types[ClusterEnvironment.__name__] += 1
	elif isinstance(plugin, LayerSync):
	if sync_batchnorm and not isinstance(plugin, TorchSyncBatchNorm):
	raise MisconfigurationException(
	f"You set `Trainer(sync_batchnorm=True)` and provided a `{plugin.__class__.__name__}`"
	" plugin, but this is not allowed. Choose one or the other."
	)
	self._layer_sync = plugin
	plugins_flags_types[TorchSyncBatchNorm.__name__] += 1
	else:
	raise MisconfigurationException(
	f"Found invalid type for plugin {plugin}. Expected one of: Precision, "
	"CheckpointIO, ClusterEnviroment, or LayerSync."
	)

[aaronwtr]

I suspect that the default behaviour of Pytorch CLI should be to interpret the plugins flag and then instantiate the appropriate corresponding object, in this case a DDPSharded type object. This does not seem to be happening

[mauvilsa]

The CLI just does what the type hint tells. The type of plugins is in part Union[PLUGIN, str]. You give a string, so the parser allows it and gives this string as value when instantiating the trainer. That logic is working correctly. If the plugins parameter doesn't really accept strings, then the type hint is wrong. That is a bug.

For the CLI to instantiate a class before giving it to the trainer, you need to give class_path and optionally init_args as explained in lightning_cli_advanced.html or in class-type-and-sub-classes.

[aaronwtr]

I did not know that, I see. I have opened an issue (here #19128) to troubleshoot if indeed the parsing is buggy for ddp_sharded. Thank you!

[wconnell]

I think the strategy arg is unique in that you can pass a string that is parsed into a class