Method test_composite_learner

axlearn/common/trainer_test.py:1050–1093

Tests composite learner with two sub learners for weight/bias respectively.

(self)

        model_cfg = DummyModel.default_config().set(dtype=jnp.float32)

        def checkpoint_if_all_evalers_run(evaler_names: Sequence[str]) -> CheckpointPolicy:
            def fn(*, step: int, evaler_summaries: dict[str, Any]):
                del step
                for evaler_name in evaler_names:
                    if evaler_summaries.get(evaler_name) is None:
                        return False
                return True

            return fn

        cfg: SpmdTrainer.Config = SpmdTrainer.default_config().set(
            name="test_trainer",
            dir=tempfile.mkdtemp(),
            mesh_axis_names=("data", "model"),
            mesh_shape=(1, 1),
            model=model_cfg,
            input=DummyInput.default_config(),
            learner=learner.Learner.default_config().set(
                optimizer=config_for_function(optimizers.sgd_optimizer).set(
                    learning_rate=0.1, decouple_weight_decay=True
                ),
            ),
            max_step=8,
            checkpointer=Checkpointer.default_config().set(
                save_policy=config_for_function(checkpoint_if_all_evalers_run).set(
                    evaler_names=["eval_every_2", "eval_every_3"]
                ),
            ),
            evalers=dict(
                eval_every_2=SpmdEvaler.default_config().set(
                    input=DummyInput.default_config().set(total_num_batches=1),
                    eval_policy=config_for_function(eval_every_n_steps_policy).set(n=2),
                ),
                eval_every_3=SpmdEvaler.default_config().set(
                    input=DummyInput.default_config().set(total_num_batches=2),
                    eval_policy=config_for_function(eval_every_n_steps_policy).set(n=3),
                ),
            ),
            save_input_iterator=save_input_iterator,
        )
        cfg.checkpointer.storage.max_concurrent_gb = max_concurrent_gb

        # Run trainer.
        trainer: SpmdTrainer = cfg.instantiate(parent=None)
        trainer.run(prng_key=jax.random.PRNGKey(123))

        assert os.path.exists(os.path.join(cfg.dir, "trainer_state_tree.txt"))