_________________ test_algorithm_resumption[1-SeqLengthWarmup] _________________

tmp_path = PosixPath('/tmp/pytest-of-root/pytest-0/test_algorithm_resumption_1_Se1')
alg_cls = <class 'composer.algorithms.seq_length_warmup.seq_length_warmup.SeqLengthWarmup'>
world_size = 1

    @pytest.mark.gpu
    @pytest.mark.parametrize('alg_cls', get_algs_with_marks())
    @pytest.mark.filterwarnings(
        'ignore:Detected call of `lr_scheduler.step()',
    )  # optimizer.step() sometimes skipped when NaN/inf on low batch size
    @pytest.mark.filterwarnings(r'ignore:.*Plan failed with a cudnnException.*:UserWarning')  # Torch 2.3 regression
    @world_size(1, 2)
    def test_algorithm_resumption(
        tmp_path: pathlib.Path,
        alg_cls: Type[Algorithm],
        world_size,
    ):
        folder1 = os.path.join(tmp_path, 'folder1')
        folder2 = os.path.join(tmp_path, 'folder2')
        os.makedirs(folder1, exist_ok=True)
        os.makedirs(folder2, exist_ok=True)
    
        model = get_alg_model(alg_cls)
        alg_kwargs = get_alg_kwargs(alg_cls)
    
        copied_model = copy.deepcopy(model)  # copy the model so the params will start from the same point
    
        if alg_cls is LayerFreezing:
            pytest.xfail('Known issues')
    
        if alg_cls in (SAM, StochasticDepth):
            pytest.xfail('Mismatch in weights when resuming from a checkpoint.')
    
        if alg_cls is GyroDropout:
            pytest.xfail('GyroDropoutLayer is not implemented in a way that allows correct resumption.')
    
        if alg_cls is SWA and world_size > 1:
            pytest.xfail('SWA is not implemented in a way that is compatible correct resumption on multiple devices.')
    
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5)
    
        shared_config = {
            'max_duration': '2ep',
            'save_filename': 'ep{epoch}-rank{rank}',
            'save_interval': '1ep',
            'train_subset_num_batches': 2,
            'precision': 'amp_bf16',
        }
        train_dataloader = get_alg_dataloader(alg_cls) if world_size == 1 else get_alg_dataloader(alg_cls, multigpu=True)
        # train model once, saving checkpoints every epoch
        trainer1 = Trainer(
            model=model,
            train_dataloader=train_dataloader,
            optimizers=optimizer,
            schedulers=scheduler,
            save_folder=folder1,
            algorithms=alg_cls(**alg_kwargs),
            **shared_config,
        )
        trainer1.fit()
    
        # create second trainer, load an intermediate checkpoint
        # and continue training
    
        optimizer = torch.optim.Adam(copied_model.parameters(), lr=0.01)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5)
    
        alg = alg_cls(**alg_kwargs)
        # SeqLengthWarmup has a call to ._activate_model() that happens on the first call to the algorithm
        # in order to get complete matching of the rng state, we have to cause that extra call to be skipped
        # when reloading.
        if alg_cls is SeqLengthWarmup:
            alg._activated = True  # type: ignore
    
        train_dataloader = get_alg_dataloader(alg_cls) if world_size == 1 else get_alg_dataloader(alg_cls, multigpu=True)
>       trainer2 = Trainer(
            model=copied_model,
            train_dataloader=train_dataloader,
            load_path=os.path.join(folder1, 'ep1-rank{rank}'),
            load_weights_only=False,
            load_strict_model_weights=False,
            optimizers=optimizer,
            schedulers=scheduler,
            save_folder=folder2,
            algorithms=alg,
            **shared_config,
        )

/composer/tests/algorithms/test_algorithm_resumption.py:91: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/composer/composer/trainer/trainer.py:1715: in __init__
    self._rng_state = checkpoint.load_checkpoint(
/composer/composer/utils/checkpoint.py:529: in load_checkpoint
    rng_state_dicts = _restore_checkpoint(
/composer/composer/utils/checkpoint.py:1006: in _restore_checkpoint
    state.load_state_dict(
/composer/composer/core/state.py:1425: in load_state_dict
    self.load_optim_state(state)
/composer/composer/core/state.py:1338: in load_optim_state
    set_optimizer_state_dict(
/composer/composer/trainer/mosaic_fsdp_utils.py:965: in set_optimizer_state_dict
    _load_optim_state_dict(model, optimizers, optim_state_dict, info)
/usr/lib/python3/dist-packages/torch/distributed/checkpoint/state_dict.py:589: in _load_optim_state_dict
    optim_state_dict = _split_optim_state_dict(model, optim, state_dict, info)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

model = HuggingFaceModel(
  (model): BertForMaskedLM(
    (bert): BertModel(
      (embeddings): BertEmbeddings(
        (word...se_affine=True)
        )
        (decoder): Linear(in_features=128, out_features=30522, bias=True)
      )
    )
  )
)
optim = Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    capturable: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: None
    initial_lr: 0.01
    lr: 0.01
    maximize: False
    weight_decay: 0
)
optim_state_dict = {'param_groups': [{'amsgrad': False, 'betas': (0.9, 0.999), 'capturable': False, 'differentiable': False, ...}], 'stat...000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00]]), 'step': tensor(2.)}, ...}}
info = _StateDictInfo(full_state_dict=False, cpu_offload=False, ignore_frozen_params=False, keep_submodule_prefixes=True, str..._prefixes=set(), handle_model=False, handle_optim=True, fsdp_context=<class 'contextlib.nullcontext'>, fsdp_modules=[])

    def _split_optim_state_dict(
        model: nn.Module,
        optim: torch.optim.Optimizer,
        optim_state_dict: OptimizerStateType,
        info: _StateDictInfo,
    ) -> OptimizerStateType:
        """
        Extract the corresponding optim state_dict from ``optim_state_dict`` for
        ``optim`` and return the result optim state_dict.
    
        Args:
            model (nn.Module): the root model.
            optim (torch.optim.Optimizer): the optimizer.
            optim_state_dict (Dict[str, ValueType]): the superset optim state_dict that
                contains the optim state_dict of ``optim``.
            info (_StateDictInfo): state dict information.
    
        Returns:
            The optim state_dict of ``optim``.
        """
    
        state: DictValueType = {}
        pg_state: ListDictValueType = []
        return_osd: OptimizerStateType = {STATE: state, PG: pg_state}
        pg_mapping: Dict[int, int] = {}
    
        for param_group in optim.param_groups:
            pg_state.append({PARAMS: []})
            for param in param_group[PARAMS]:
                for fqn in info.fqn_param_mapping[param]:
                    params = pg_state[-1][PARAMS]
                    assert isinstance(params, list)
                    params.append(fqn)
                    if param.requires_grad:
>                       state[fqn] = cast(DictValueType, optim_state_dict[STATE])[fqn]
E                       KeyError: 'model.cls.predictions.decoder.weight'

/usr/lib/python3/dist-packages/torch/distributed/checkpoint/state_dict.py:559: KeyError
------------------------------ Captured log setup ------------------------------

        fqn_param_mapping: Dict[
            Union[str, torch.Tensor], Union[Set[str], torch.Tensor],
        ] = {}
        for name, param in chain(model.named_parameters(), model.named_buffers()):
            fqns = _get_fqns(model, name)
            fqn_param_mapping[param] = fqns
            for fqn in fqns:
                fqn_param_mapping[fqn] = param

        all_fqns = set()
        for name, _ in _iterate_valid_model_state(model):
            fqns = _get_fqns(model, name)
            for fqn in fqns:
                all_fqns.add(fqn)