This will make output_args required (right now it is optional), otherwise initialization will fail right? To have a better transition, can we still run self.submod() with the input args and leave the output_args to be false. This will fail if the model is not on the same device as the input, but that is okay and we can say we expect them to be both on meta device.

you're right that this makes output_args required. I like your suggestion. If you're willing to let it be an error if users pass real inputs, I would propose to assert the module and inputs are on the same device, and then compute 'outputs_meta' based on inputs.

I was weighing whether to make it a requirement to pass input/module on meta device, but I suppose that is too restrictive. I do want to ensure that outputs_meta is stored on meta to avoid wasting memory. So maybe if the user passes cuda model/inputs I will convert the output to meta. Wdyt?

ok, i added output-shape inference back in, but i made it refuse to do inference on non-meta device, so some user code might still have to switch their inputs to meta. is this ok, or should i make it work for cuda?

It probably might be easier to just let it work for cuda if the users already have it on cuda. Then you wouldn't have to fix all the tests as you mentioned above. I think the main fix of the PR can just be removing the .to(device) in init which you did, and with that removal now it is up to the users to make sure that model(input_args) works if they pass it in.

Hmm, can I say that I prefer the previous version better? i.e. removal.
(See also my comments about on relaxing device constraint, which requires removal of dryrun.)

Hmm, I'm in the other boat, I think keeping the dry run makes more sense (as a temporary measure) until we get the lazy init working, otherwise we have to update all the tests to pass in both input_args and output_args and also requires users to do so as well.

Yeah, that's hard. Maybe at least add with torch.no_grad()?
I am actually not sure what side effect running a module has. Saving grad context, marking requires_grad?

Also, in forward_maybe_with_... we have special logic for DDP/FSDP modules, would this dryrun's escape be an issue? (Or are we just lucky?)

i'm going to leave this for later. its not any worse with this PR and we plan to remove it when adding lazy shape inference. sound ok?

Yep, sounds good.

hm, should I move this logic inside _step_microbatches? it seems some of our tests are calling _step_microbatches. Do we allow this or do we require calling step()?

if we don't require calling step(), then should we also move this code inside _step_microbatches?

        self._stage.clear_runtime_states()

why don't i just do this init inside stage._forward_one_chunk...?

Yeah. That would be a good place.
We may have to ask CUDA Graph to come capture after the first step(...) though.
But that may not be a big deal? And the current change moved the prepare from init to step anyway.

annoying: Stage class doesn't know "n_microbatches", which may have been intentional?

I am thinking about whether it is bad to let the schedule 'register' the number of microbatches during Schedule.init and then stage can use this value later when it performs initialization inside forward_one_chunk / backward_one_chunk?

this assertion breaks tests. i guess itll break some existing usages. Is it a good thing to do to prevent bad practices in the future, or should i relax this and make it work if model+inputs are on cuda?

I think we should relax this. "meta" device is still a high-end thing for most users.
And, is nn.Module.device a real thing?

hmm, its not always possible to assert this- maybe there is another way to reliably check the device of the module?

AttributeError: 'FSDPSequential' object has no attribute 'device'

yea, maybe i should just drop the checks and perform the shape inference blindly, but then do a .to(meta) on the outputs. that should make things on par for what howard mentioned, and then we can delete it once we add lazy inference

FSDPSequential is a superset of 'nn.Sequential'. It means vanilla nn.Sequential does not have device (see below for repro). I guess what we really want to check whether model.parameters() are on meta? that should cover all cases

>>> import torch
>>> model = torch.nn.Sequential(torch.nn.Linear(2,2))
>>> model.device
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/data/users/weif/pytorch/torch/nn/modules/module.py", line 1931, in __getattr__
    raise AttributeError(
AttributeError: 'Sequential' object has no attribute 'device'

I think we should relax this. "meta" device is still a high-end thing for most users.
And, is nn.Module.device a real thing?

Hmm, can I say that I prefer the previous version better? i.e. removal.
(See also my comments about on relaxing device constraint, which requires removal of dryrun.)

Oh then should we disable validation logic (which an add-on protection) till it supports the multi-output case? Multi-output is pretty common. This assert may break a few tracer cases.

if this is a new regression, i can remove the assert, but i thought the assert would just make another error more explicit. The old code assumes 'output' is a single value and tries to validate against it. Lets see if CI fails and if not, make another PR to fix validation?

I see. Yep, if it is not a new assert, then let's see what CI says. :)

Yep, sounds good.

I see. Yep, if it is not a new assert, then let's see what CI says. :)