Demonstrating Monty's Capabilities: Configs Code-Share

[scottcanoe]

This is largely unchanged from the code-share in the last repo except for a couple of pretraining parameters. I expect this to undergo fairly significant changes and additions in the near future as we add more experiments.

Some thoughts:

• It would probably be best to do this all in monty_lab. Right now, we are blocked from doing that because custom config dataclasses (like those for 9 and 10 LM configs) must be in tbp.monty in order to load model.pt files. I wonder if this might be a pain point for users wanting to use the config dataclass pattern.
• A workaround is to use functions that return instances of core tbp.monty config dataclasses, and the functions do the work of customizing config parameters. I adopted that pattern in the dmc_pretraining_experiments.py and dmc_eval_experiments.py before realizing it might be a solution to the above problem.
• It's on my list to write functions to generate configs for multi-LM monty and dataset configs. Once that is done, we should be able to move to monty_lab without issue. It'll also help switch to different SM/LM setups since 1, 2, 5, 9, 10 is kind of arbitrary and non-ideal. We talked about doing something like 1, 2, 4, 8, 16 SMs/LMs.
• We might want to repackage this code in the future where different sets of experiments get broken into different files. eval_experiments.py is already getting pretty long, and it'll get quite a bit longer when we add the rapid learning experiments. In this case, it might make sense to create a module like dmc_common.py which contains parameters and reusable config-getter functions, rotations, etc.

[vkakerbeck]

Thanks for sharing! Is there a reason you are adding these to the tbp.monty repository instead of the monty_lab repo? Since these experiments are npot part of our benchmark test suite I would not add them here in the benchmarks folder. Can we put these configs into this folder https://github.com/thousandbrainsproject/monty_lab/tree/main/monty_capabilities_analysis ? Similar to how we have a bunch of configs here: https://github.com/thousandbrainsproject/monty_lab/tree/main/experiments/configs

[hlee9212]

It would probably be best to do this all in monty_lab. Right now, we are blocked from doing that because custom config dataclasses (like those for 9 and 10 LM configs) must be in tbp.monty in order to load model.pt files. I wonder if this might be a pain point for users wanting to use the config dataclass pattern.

Slightly tangential but I'm also testing out having all configs in monty_lab and will see if I can take this Draft PR and create a new one in monty_lab repo.

[vkakerbeck]

Maybe you can give me some more details some time of why we can't define new dataclasses in Monty lab? In general, if there is no way to do this, we could make this PR here only contain the new dataclasses (and actually merge the PR) and then have the configs in monty_lab

@tristanls just as an fyi of another current pain point with configs.

[scottcanoe]

Well... I ran into this problem with loading models before when dataclasses weren't in tbp.monty, but I just successfully loaded a 9-LM model.pt while on main and had no problem.

So what I think happened was this: about a month ago, I was saving checkpoints during pretraining so I could keep track of how many points were added per epoch. This saved config.pt files, and loading those was throwing errors because config.pt had pickled everything in a monty experiment. Since the custom dataclasses had to be visible for unpickling, loading config.pt without those classes in tbp.monty threw an error. I didn't keep detailed notes, but I remember there was a class discovery mechanism in torch that wasn't seeing classes defined outside tbp.monty.

I had written a custom logger for saving checkpoints, so maybe it was an issue on my end. Like maybe I needed to be dict-ifying the experiment config before saving it.

I think the good news is that I was wrong about not being able to load model.pt files. On the other hand, I'll be revisiting checkpointing soon for the rapid learning experiments. While I'm at it, I'll try to nail down what the issue was exactly.

[vkakerbeck]

Okay interesting, thanks for the extra infos! I'll be curious to hear what you find out.

[hlee9212]

Re: "dict-ifying" experimental config -- that's the top thing I want, haha. Instead of saving the dataclass itself, to save a serialized version.

[nielsleadholm]

Thanks for sharing those details Scott.

Yeah in terms of making complex data-classes serializable, I'm curious what the best options are, e.g. it might just be worth implementing some custom methods since we don't have that many different types of objects. Something we can continue to keep in mind as part of the RFC in PR 46

[hlee9212]

I learned that pydantic package can do this automatically with its built-in model_dump() (saves as dictionary) model_dump_json() (saves as JSON) 😄 (link: https://docs.pydantic.dev/latest/concepts/serialization/)

[nielsleadholm]

Ah ok interesting! Yeah this description certainly sounds interesting (form the main 'pydantic' page)

"Models share many similarities with Python's dataclasses, but have been designed with some subtle-yet-important differences that streamline certain workflows related to validation, serialization, and JSON schema generation. You can find more discussion of this in the Dataclasses section of the docs."

[hlee9212]

Hi @scottcanoe, @vkakerbeck, @nielsleadholm - I moved the configs from this PR to monty_lab here. I checked with pretrain_dist_agent_1lm that the config loads and starts training.

I think this PR can now be closed.

[nielsleadholm]

Thanks @hlee9212 !