Error when initializing DistributedTrainer

[manuelbrack]

Doc string of DistributedTrainer does not align with init behavior.

According to the docstring the config_or_config_file can be either a Config object or a path to a config.

config_or_config_file: Either a Config object or a path to a YAML file containing the config.

nanotron/src/nanotron/trainer.py

Lines 119 to 135 in d57c2ea

	class DistributedTrainer:
	def __init__(
	self,
	config_or_config_file: Union[Config, str],
	config_class: Type[Config] = Config,
	model_config_class: Optional[Type] = None,
	model_class: Type[NanotronModel] = None,
	):
	"""
	Nanotron's distributed trainer.
	Args:
	config_or_config_file: Either a `Config` object or a path to a YAML file containing the config.
	config_class: The `Config` class to use.
	model_config_class: The `ModelConfig` class to use (for example `LlamaConfig`). Defaults to `None` which will use the model config class defined in the config.
	model_class: The `NanotronModel` class to use (for example `LlamaForTraining`). Defaults to `None` which will use the model class defined in the config.
	"""

However, the current initialization using get_config_from_file only allows for a file path and not a config object.

nanotron/src/nanotron/trainer.py

Lines 138 to 140 in d57c2ea

	self.config = get_config_from_file(
	config_or_config_file, config_class=config_class, model_config_class=model_config_class
	)

The easiest fix would probably be to run a check if config_or_config_file is already a Config object and set self.config = config_or_config_file in that case instead of calling get_config_from_file .