Support memory_format on to()

[tfogal]

🚀 Feature

to(memory_format=something) is part of the MegatronImagen model in NeMo.

Ideally, this would work:

$ git diff .
diff --git a/nemo/collections/multimodal/models/text_to_image/imagen/imagen.py b/nemo/collections/multimodal/models/text_to_image/imen/imagen.py
index 4fa6cd230..2cf7a8ffa 100644
--- a/nemo/collections/multimodal/models/text_to_image/imagen/imagen.py
+++ b/nemo/collections/multimodal/models/text_to_image/imagen/imagen.py
@@ -31,6 +31,7 @@ from nemo.collections.nlp.modules.common.megatron.module import Float16Module
 from nemo.collections.nlp.parts.utils_funcs import get_last_rank
 from nemo.core.classes.common import Serialization
 from nemo.utils import logging
+import thunder
 
 try:
     from apex import amp
@@ -190,6 +191,7 @@ class MegatronImagen(MegatronBaseModel):
         self.megatron_amp_O2 = cfg.get('megatron_amp_O2', False)
 
         self.model = self.model_provider_func()
+        self.model = thunder.jit(self.model)
 
         if self.trainer.precision in ['bf16', 'bf16-mixed']:
             self.autocast_dtype = torch.bfloat16

Motivation

Trying to evaluate NeMo models in thunder and expand our model support there. Megatron-based models appear to be widely used.

Alternatives

I wonder if we could temporarily just accept the keyword without actually doing anything about it. I imagine that would be very slow, but it might allow us to get models like this one into thunder more easily.

I'll start trying to convert smaller parts of the model next.

Additional context

Model in question:

https://github.com/NVIDIA/NeMo/blob/23baa48e441ecb6cc6b49c23bf8cfc076db38bdc/nemo/collections/multimodal/models/text_to_image/imagen/imagen.py#L175

I think the to that is failing for me
is actually this line:
https://github.com/NVIDIA/NeMo/blob/23baa48e441ecb6cc6b49c23bf8cfc076db38bdc/nemo/collections/multimodal/models/text_to_image/imagen/imagen.py#L135

Model test:
log.txt

[jjsjann123]

I think we just didn't add the memory_format arg in our thunder/torch/__init__.py

we should be easily mapping it to stride_order prim.

I'll take a stab on this one.