About implementation of CompileGraphModel

[hyuntae-cho]

Hi,

I've been exploring how torch.compile is applied to NequIP and had a question regarding the ListInputOutputStateDictWrapper in nequip/nn/compile.py.

It looks like the forward function copies the parameters and buffers from the original model. However, when I insert the following line:

print(set([v.data_ptr() for k,v in self._compiled_model[0].named_parameters()]) == set([v.data_ptr() for k,v in self.model.named_parameters()]))

right after this block:

nequip/nequip/nn/compile.py

Lines 174 to 180 in 60eae37

	self._compiled_model = (
	torch.compile(
	fx_model,
	dynamic=True,
	fullgraph=True,
	),
	)

the printed result is True, indicating that the compiled model shares the same parameter tensors (by memory address) as the original model.

I even tried removing the parameter copying code:

nequip/nequip/nn/compile.py

Lines 87 to 91 in 60eae37

	with torch.no_grad():
	for name, param in self.model.named_parameters():
	param.data.copy_(state_dict[name])
	for name, buffer in self.model.named_buffers():
	buffer.data.copy_(state_dict[name])

and still obtained the same results.

It seems this doesn't impact computational flow or cost, so I’m curious—was there a specific reason for implementing the wrapper with explicit copying of parameters and buffers? Is it for compatibility with an older PyTorch version or to cover edge cases?

Thanks in advance for your insight!

Best,
Hyuntae Cho

[Linux-cpp-lisp]

Hi Hyuntae,

Thanks for the interesting question, I'm not sure off the top of my head...

I'm curious, are you implementing something specific in or with nequip that made you need to modify or look at the compilation code? (Please also always feel free to email us at [email protected].)

[cw-tan]

Hi @hyuntae-cho ,

The goal of the nn.Module was just to functionalize the model such that it looks like

func(**true_inputs, **trainable_params, **buffers) -> **true_outputs

for the make_fx procedure (which also "flattens" out the autograd for forces). Then, it's seems a bit cleaner to apply torch.compile on it and get compiled autograd for training the trainable_params. It is likely that one might set it up differently and have it all work, but it felt clearer to literally functionalize the model (as in the signature above) before compiling.