This error occurs when I tried to run cifar10 example on GPU (TITAN XP)

Error log

Traceback (most recent call last):
  File "vmoe/main.py", line 71, in <module>
    app.run(main)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/absl/app.py", line 303, in run
    _run_main(main, args)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "vmoe/main.py", line 64, in main
    trainer.train_and_evaluate(config=FLAGS.config, workdir=FLAGS.workdir)
  File "/mnt/cephfs/home/cascol/working-directory/MoE-in-MoE/vmoe/vmoe/train/trainer.py", line 434, in train_and_evaluate
    tree=jax.eval_shape(train_state_initialize_fn, train_state_rngs),
  File "/mnt/cephfs/home/cascol/working-directory/MoE-in-MoE/vmoe/vmoe/train/trainer.py", line 203, in initialize
    variables = model.init(rngs, inputs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/flax/linen/module.py", line 1124, in init
    method=method, mutable=mutable, **kwargs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/flax/linen/module.py", line 1092, in init_with_output
    {}, *args, rngs=rngs, method=method, mutable=mutable, **kwargs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/flax/linen/module.py", line 1061, in apply
    )(variables, *args, **kwargs, rngs=rngs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/flax/core/scope.py", line 706, in wrapper
    with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/flax/core/scope.py", line 684, in bind
    'rngs should be a dictionary mapping strings to `jax.PRNGKey`.')
jax._src.traceback_util.UnfilteredStackTrace: flax.errors.InvalidRngError: rngs should be a dictionary mapping strings to `jax.PRNGKey`. (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.InvalidRngError)

My Enviroment that may concern

cudatoolkit               11.1.1               h6406543_8    conda-forge
flax                      0.3.6                    pypi_0    pypi
jax                       0.2.27                   pypi_0    pypi
jaxlib                    0.1.75+cuda11.cudnn805          pypi_0    pypi
optax                     0.1.0                    pypi_0    pypi
pytorch                   1.8.0           py3.7_cuda11.1_cudnn8.0.5_0    pytorch

I tried tracking the error and found weird bug (the image on top is the result while setting up training, the image at the bottom is what I think the result should look like)

Hello! Can I know if I need to do any changes to the scripts before using the checkpoints for evaluation? I downloaded the vmoe_b16_imagenet21k_randaug_strong_ft_cifar10 checkpoint files (both .index and .data-00000-of-00001 ) and named them ckpt_1.index and ckpt_1.data-00000-of-00001, respectively. Also, for running the script on a single partition I changed the "num_expert_partitions" of config dict to 1. With the above changes, when I try to run the script on google colab using below command: python vmoe/main.py --workdir=./vmoe/saved_checkpoints --config=vmoe/configs/vmoe_paper/vmoe_b16_imagenet21k_randaug_strong_ft_cifar10.py (where saved_checkpoints is the directory where the checkpoint files are present) I get the following error: ValueError: Missing field step in state dict while restoring an instance of TrainState Any help is appreciated!!

Complete error stack:

Traceback (most recent call last):
  File "vmoe/main.py", line 77, in <module>
    app.run(main)
  File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "vmoe/main.py", line 70, in main
    trainer.train_and_evaluate(config=FLAGS.config, workdir=FLAGS.workdir)
  File "./vmoe/train/trainer.py", line 527, in train_and_evaluate
    return _train_and_evaluate(config, workdir, mesh)
  File "./vmoe/train/trainer.py", line 575, in _train_and_evaluate
    thread_pool=ThreadPool())
  File "./vmoe/train/trainer.py", line 291, in restore_or_create_train_state
    thread_pool=thread_pool)
  File "./vmoe/checkpoints/partitioned.py", line 83, in restore_checkpoint
    'index': tree if tree is not None else axis_resources,
  File "./vmoe/checkpoints/base.py", line 137, in restore_checkpoint
    return serialization.from_bytes(tree, checkpoint_contents)
  File "./vmoe/checkpoints/serialization.py", line 71, in from_bytes
    return from_state_dict(target, state_dict)
  File "/usr/local/lib/python3.7/dist-packages/flax/serialization.py", line 65, in from_state_dict
    return ty_from_state_dict(target, state)
  File "/usr/local/lib/python3.7/dist-packages/flax/serialization.py", line 128, in _restore_dict
    for key, value in xs.items()}
  File "/usr/local/lib/python3.7/dist-packages/flax/serialization.py", line 128, in <dictcomp>
    for key, value in xs.items()}
  File "/usr/local/lib/python3.7/dist-packages/flax/serialization.py", line 65, in from_state_dict
    return ty_from_state_dict(target, state)
  File "/usr/local/lib/python3.7/dist-packages/flax/struct.py", line 146, in from_state_dict
    raise ValueError(f'Missing field {name} in state dict while restoring'
ValueError: Missing field step in state dict while restoring an instance of TrainState

The VMOE model is great. I wonder when the pre-trained config would be made available? The pre-trained models are all based on the ImageNet-21k data, which makes sense and great. But could you guys also release the pre-trained models based on ImageNet-1k? That could be very helpful for benchmark and comparisons.

Thanks and looking forward to your responses.

[NumPy] Remove references to deprecated NumPy type aliases.

This change replaces references to a number of deprecated NumPy type aliases (np.bool, np.int, np.float, np.complex, np.object, np.str) with their recommended replacement (bool, int, float, complex, object, str).

NumPy 1.24 drops the deprecated aliases, so we must remove uses before updating NumPy.

[NumPy] Remove references to deprecated NumPy type aliases.

This change replaces references to a number of deprecated NumPy type aliases (np.bool, np.int, np.float, np.complex, np.object, np.str) with their recommended replacement (bool, int, float, complex, object, str).

NumPy 1.24 drops the deprecated aliases, so we must remove uses before updating NumPy.

[JAX] Avoid implicit references to jax._src.

An upcoming change to JAX means that it will no longer export jax._src by default. It still works (for now) to import modules from jax._src and refer to those, but they will not be present in JAX's namespace except via an explicit import.

Note that any use of jax._src is a use of a JAX-private API. Please use public APIs instead. This change mostly does not yet switch users onto public APIs.

Inject current learning rate to the optimizer state. Allow plotting the norm of any array in the train state.

We can now plot the norm of current parameter values, the learning rate or any other array part of the optimizer state:

config.plot_norm_train_state_patterns = [
  # norm of the kernel params in MoE layers.
  'params/.*/moe/mlp/.*/kernel',
  # norm (=value) of current learning rate.
  'opt_state/.*/hyperparameter/learning_rate',
  # norm of 1st order moments of the kernel gradients in MoE layers (i.e. Adam inner state).
  'opt_state/.*/mu/.*/moe/mlp/.*/kernel',
]

Hi，I want to ask a question, even if each expert sets the capacity, will the router algorithm send all the tokens to one or two experts during training?

This may sound trivial. But I'd like to have some idea on how use the .data and .index file from gs://vmoe_checkpoints. Does this repo provide any scripts or code to do some evaluation/inference on those models? Any help would be appreciated.