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

[yhtang]

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`.

[Ethanlm]

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

[apghml]

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)

[yhtang]

That should work just fine, like how we captured things in the idiomatic way of defining decorators.
Here is a simple test that captures the pattern:

def f():
  x = 1
  return lambda: print(x)

g = f()
g()

which correctly prints 1 when we called g().

[apghml]

Maybe I missed something, but why doesn't the code in the PR exhibit the same problem as the following code:

fns = []
for i in range(3):
    fns.append(lambda: i)
print( fns[0]() ) # Incorrectly prints 2 instead of 0.

I agree your snippet doesn't have this issue because the value being closed over is not modified.

[yhtang]

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

This is only triggered if using one of the *-single-host configs, e.g. fuji-3B-v3-flash-single-host

[Ethanlm]

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

This is only triggered if using one of the *-single-host configs, e.g. fuji-3B-v3-flash-single-host

Yeah I have been using single-host config for my testing on TPU recently.

[ruomingp]

Will defer to @apghml for approval.