Unable to find the compile_stage library

[dheerj188]

How can i import the compile_stage library to initialize my stage? it is not found after installation.

### Tasks

[dheerj188]

The examples demonstrated indicate that compile_stage resides in pippy or pippy.compile. I could not find any one of them.

[kwen2501]

Hi, the compile_stage API is deprecated. Some up-to-date examples include: examples/basic/example.py examples/huggingface examples/llama

[dheerj188]

So what's the guidelines for training? should i create remote optimizers for each stage and call the step method?

[kwen2501]

Here is a simple training + optimizer example: https://github.com/pytorch/PiPPy/blob/main/test/test_optim.py

For backward, you can pass a loss function to PipelineSchedule:

    # Attach to a schedule
    schedule = PipelineScheduleGPipe(stage, args.chunks, loss_fn=loss_fn)

At run time, you can pass the target to the last rank:

        elif args.rank == args.world_size - 1:
            losses = []
            out = schedule.step(target=target, losses=losses)

(losses is an optional pass-in container for PiPPy to return the per-microbatch loss.)

For optimizer, you can attach a local optimizer per stage:

    # Create an optimizer for stage submodule's parameters
    optimizer = optim.SGD(stage.submod.parameters(), lr=1e-3, momentum=0.9)

In the training step, you can call it like a normal optimizer:

        # Take an optimization step
        optimizer.step()