Skip to content

Use lambda instead of functools.partial to create single-host config maker #1152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

yhtang
Copy link

@yhtang yhtang commented May 2, 2025

Fixes #1151

An alternative solution is to use functools.wraps but that requires a custom config validator:

axlearn.common.config.InvalidConfigValueError: Invalid config value type <class 'functools.partial'> for value "functools.partial(<function trainer_configs.<locals>.make_single_host_config at 0x722900936160>, 'fuji-3B-v3-flash')". Consider registering a custom validator with `register_validator`.

@yhtang yhtang requested review from ruomingp, markblee and a team as code owners May 2, 2025 21:10
@Ethanlm
Copy link
Contributor

Ethanlm commented May 2, 2025

I have not encountered this issue. How can I reproduce it?

@@ -1011,7 +1011,6 @@ def make_single_host_config(base_config_name: str) -> SpmdTrainer.Config:
return cfg

# Make single-host config
make_single_host_config_func = functools.partial(make_single_host_config, config_name)
config_map[f"{config_name}-single-host"] = make_single_host_config_func
config_map[f"{config_name}-single-host"] = lambda: make_single_host_config(config_name)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIUC, this won't work as is due to how closures work in Python?

Suggested change
config_map[f"{config_name}-single-host"] = lambda: make_single_host_config(config_name)
config_map[f"{config_name}-single-host"] = lambda config_name=config_name: make_single_host_config(config_name)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Single-host config maker rejected by config_class_for_function
3 participants