Is there a way to export a pipeline stage?

[nrs-status]

I'd like to load a pipeline stage into a VM with low disk storage. Is there a way to export a pipeline stage and have it run independently from the file responsible for wrapping the model into Pipe?

[kwen2501]

Hi, in the topmost of main, we added a new API:

pipe = pippy.pipeline(model, ...)
stage_module = pipe.get_stage_module(stage_idx)

I think you can export stage_module in a couple ways:

Method 1. torch.export

exported_program = torch.export.export(stage_module, ...)

(nit: you may need an example input for stage_module for the ... part for export to work)

I think torch.export also has save and load methods, e.g.:

torch.export.save(exported_program, PATH)

Method 2. torch.save

torch.save(stage_module.state_dict(), PATH)

On the new VM with limited memory, use meta device to create the stage_module:

with torch.device("meta"):
    model = ModelConstructor(...)

pipe = pippy.pipeline(model, ...)

# stage_module will be on meta device as well
stage_module = pipe.get_stage_module(stage_idx)

# Load real weights into it
stage_module.load_state_dict(torch.load(PATH))

[nrs-status]

Thanks a lot! Thanks for the extensive comments and instructions, highly appreciated. Will be using this to do some tests on a distributed setup.