DeepSpeed - Q3 Roadmap

[PKUWZP]

This is the roadmap for DeepSpeed Q3 2026. Feedback is welcome — please leave comments on this issue or join the #2026q3-roadmap channel on the DeepSpeed Slack.

Some of the items are the extension of the Q2 roadmap: #7861

New feature and enhancement

Enhance Expert-Parallelism (EP) support

• AutoEP Support: AutoEP enables Expert Parallelism (EP) for major Mixture-of-Experts (MoE) models out of the box, eliminating the need for users to write model-specific parallelization code. By automatically distributing expert layers across devices, AutoEP allows users to scale MoE training with minimal configuration changes.
• Combining AutoEP with AutoTP
• Extend AutoTP capabilities by integrating Liger Kernel (sharding LM heads and adopting Online Softmax)
• EP performance optimization and benchmarking

New Accelerators Support (Q3)

• DeepSpeed TPU accelerators support
• Performance Optimization on Emerging AI Accelerators (e.g. Biren AI Accelerators)

On-Policy Distillation Trainer Support (Q3)

• Systems Design, prototyping and benchmarking (reference: https://arxiv.org/abs/2604.14084)

DeepCompile efficiency and robustness improvements

• Formal pass contracts and validation of optimization passes: Add lightweight optimization pass contracts for automatic compatibility validation and ordering.
• Composable AutoTP / AutoEP / SP optimization passes: Implement AutoTP and AutoEP at a compiler level and integrate them with AutoSP.
• Enhance optimization passes: Add some optimization passes including activation offloading.

Pipeline parallelism with Ray

• Ray-backed pipeline stage placement and execution: Prototype pipeline stage placement and execution using Ray actor groups for a representative transformer workload, enabling heterogeneous resource allocation across pipeline stages. (Will be a generalized implementation of this).
• Create some examples: Provide example implementations of Ray-backed pipeline parallelism for representative transformer workloads, demonstrating the benefits of heterogeneous resource allocation.

Tuning guide with benchmarking results

• Recommended configurations for representative models: Provide practical configuration guidance for representative dense and MoE models across several common GPU setups.
• Benchmark-backed guidance: Run targeted benchmarks to support the recommended configurations and document the observed throughput, memory usage, and limitations.

Stability (Q3)

• Performance regression test
• Enable nightly full test
  • CUDA
  • AMD
  • Intel XPU
  • Intel Gaudi
  • NPU

[phoeenniixx]

Hi @PKUWZP , I would like to contribute to the design of On-Policy Distillation. I am new to this community, so before going forward with anything i would like to know from the maintainers if they think this is a good design.

From what i can undersatnd we can use Layer reduction with OPD (or just OPD).
What I mean by this is:

• student is either a smaller model (like in the paper there is already a chosen teacher - student pair), or
• We allow the user to create their own "student" if they dont have any pair chosen by layer reduction (or any other model compression method like pruning?)

Now we have a teacher - student pair.

Configs

Here we can either of two things:

1. if breaking changes are expected

We introduce a new distillation block inside compression_training. This allows the user to decide if they want to perform the current layer-reduction with KD or OPD.

{
  "compression_training": {
    "distillation": {
      "type": "on_policy", // options: "off_policy", "on_policy"
      "teacher_model_name_or_path": "Qwen/Qwen3-8B",
      "kd_beta_ce": 1.0,

      "student_init_method": "layer_reductino", 
      // Options: "pretrained_checkpoint", "layer_reduction", "pruning", "quantization" etc 

      "pretrained_checkpoint_config": {
        "model_name_or_path": "Qwen/Qwen3-4B"
      }, // if we use pretrained_checkpoint for student_init_method

      "layer_reduction_config": {
        "module_name_prefix": "model.layers",
        "teacher_layer": [0, 2, 4, 6, 8, 10, 12, 14],
        "other_module_name": ["model.embed_tokens", "model.norm", "lm_head"]
      }, // if we use layer_reduction for student_init_method

      "warmup": {
        "min_steps": 3000,
        "max_steps": 10000,
       ...
      }, // we might need a warmup config to realign student's layers 
      // in case we are not using a pretrained-checkpoint as student_init_method 


      "on_policy_config": {
        "rollouts_per_prompt": 16,
        "max_prompt_length": 2048,
        "max_response_length": 8192,
        "temperature": 1.0,
        "top_p": 1.0,
        "loss_chunk_size": 512,
        "token_selection": {
          "method": "soft_or", // options: "entropy_only", "q3_only" , "soft_or"
          "retention_ratio": 0.5,
          "entropy_clip_percentile": 0.98
        }
      }
    }
  }
}

2. If we dont want to make any breaking changes

We allow the layer reduction with KD same as this is now - meaning something like this (see below) works as it is.

"compression_training": {
    "layer_reduction": {
      "enabled": true,
      "keep_number_layer": 5,
      "module_name_prefix": "bert.encoder.layer",
      "teacher_layer": [
        2,
        4,
        6,
        8,
        10
      ],
      "other_module_name": [
        "bert.pooler",
        "bert.embeddings",
        "classifier"
      ]
    }
  }

And we add a on_policy_distillation block as an extension (or a new feat) in here.

{
  "compression_training": {
    "on_policy_distillation": {
      "teacher_model_name_or_path": "Qwen/Qwen3-8B",
      "kd_beta_ce": 1.0,

      // student is always a pre-trained config. 
      // Although we can allow the user to add a model compression method block before on_policy_distillation
      // that defines the student and we might need a warmup block here and not a pretrained_checkpoint_config
      
      "pretrained_checkpoint_config": {
        "model_name_or_path": "Qwen/Qwen3-4B"
      }, 

      "on_policy_config": {
        "rollouts_per_prompt": 16,
        "max_prompt_length": 2048,
        "max_response_length": 8192,
        "temperature": 1.0,
        "top_p": 1.0,
        "loss_chunk_size": 512,
        "token_selection": {
          "method": "soft_or", // options: "entropy_only", "q3_only" , "soft_or"
          "retention_ratio": 0.5,
          "entropy_clip_percentile": 0.98
        }
      }
    }
  }
}

Student Initialization

We initialise the student using:

1. the pretrained_checkpoint_config (if defined) or
2. the user uses existing model compression techniques to create a student

In the 2. case we might want a warmup block as shown in the above configs to allow the layers to realign its weights before going to the OPD process to prevent any rubbish token generation by the student during the forward loop

OPD

Once the student is initialized (by any of the paths above), we use a separate OPD flow. Because OPD requires the student to generate text, we introduce a specialized training loop wrapper.
we use an OPDTrainer wrapper that performs the standard OPD flow from the paper (student generates a rollout, teacher scores it, top-K token selection, masked KL loss).

Vignettes

Here is how the user would interact with this design in their training script:

Using configs for everything

import deepspeed
from deepspeed.compression.opd import OPDTrainer

ds_config = "ds_config.json"

opd_trainer = OPDTrainer.from_config(ds_config) 
# here we can also allow the user to pass student and teacher directly and not to the config


opd_trainer.train()

Allow passing different student/teacher model to the trainer for flexibility

opd_trainer = OPDTrainer(
    student_model=student_model,
    teacher_model=teacher_model,
    ds_config="ds_config.json",
    train_dataset=dataset
)
opd_trainer.train()

Here, we in case we allow the current model compression techniques to be defined as student and we need a warmup loop - train should handle that after looking at the configs

Or do you prefer that we dont have a trainer and the loop is written by the user themselves?

[PKUWZP]

Hi @PKUWZP , I would like to contribute to the design of On-Policy Distillation. I am new to this community, so before going forward with anything i would like to know from the maintainers if they think this is a good design.

From what i can undersatnd we can use Layer reduction with OPD (or just OPD). What I mean by this is:

• student is either a smaller model (like in the paper there is already a chosen teacher - student pair), or
• We allow the user to create their own "student" if they dont have any pair chosen by layer reduction (or any other model compression method like pruning?)

Now we have a teacher - student pair.

Configs

Here we can either of two things:

1. if breaking changes are expected

We introduce a new distillation block inside compression_training. This allows the user to decide if they want to perform the current layer-reduction with KD or OPD.

{
"compression_training": {
"distillation": {
"type": "on_policy", // options: "off_policy", "on_policy"
"teacher_model_name_or_path": "Qwen/Qwen3-8B",
"kd_beta_ce": 1.0,

  "student_init_method": "layer_reductino", 
  // Options: "pretrained_checkpoint", "layer_reduction", "pruning", "quantization" etc 

  "pretrained_checkpoint_config": {
    "model_name_or_path": "Qwen/Qwen3-4B"
  }, // if we use pretrained_checkpoint for student_init_method

  "layer_reduction_config": {
    "module_name_prefix": "model.layers",
    "teacher_layer": [0, 2, 4, 6, 8, 10, 12, 14],
    "other_module_name": ["model.embed_tokens", "model.norm", "lm_head"]
  }, // if we use layer_reduction for student_init_method

  "warmup": {
    "min_steps": 3000,
    "max_steps": 10000,
   ...
  }, // we might need a warmup config to realign student's layers 
  // in case we are not using a pretrained-checkpoint as student_init_method 


  "on_policy_config": {
    "rollouts_per_prompt": 16,
    "max_prompt_length": 2048,
    "max_response_length": 8192,
    "temperature": 1.0,
    "top_p": 1.0,
    "loss_chunk_size": 512,
    "token_selection": {
      "method": "soft_or", // options: "entropy_only", "q3_only" , "soft_or"
      "retention_ratio": 0.5,
      "entropy_clip_percentile": 0.98
    }
  }
}

}
}

2. If we dont want to make any breaking changes

We allow the layer reduction with KD same as this is now - meaning something like this (see below) works as it is.

"compression_training": {
"layer_reduction": {
"enabled": true,
"keep_number_layer": 5,
"module_name_prefix": "bert.encoder.layer",
"teacher_layer": [
2,
4,
6,
8,
10
],
"other_module_name": [
"bert.pooler",
"bert.embeddings",
"classifier"
]
}
}
And we add a on_policy_distillation block as an extension (or a new feat) in here.

{
"compression_training": {
"on_policy_distillation": {
"teacher_model_name_or_path": "Qwen/Qwen3-8B",
"kd_beta_ce": 1.0,

  // student is always a pre-trained config. 
  // Although we can allow the user to add a model compression method block before on_policy_distillation
  // that defines the student and we might need a warmup block here and not a pretrained_checkpoint_config
  
  "pretrained_checkpoint_config": {
    "model_name_or_path": "Qwen/Qwen3-4B"
  }, 

  "on_policy_config": {
    "rollouts_per_prompt": 16,
    "max_prompt_length": 2048,
    "max_response_length": 8192,
    "temperature": 1.0,
    "top_p": 1.0,
    "loss_chunk_size": 512,
    "token_selection": {
      "method": "soft_or", // options: "entropy_only", "q3_only" , "soft_or"
      "retention_ratio": 0.5,
      "entropy_clip_percentile": 0.98
    }
  }
}

}
}

Student Initialization

We initialise the student using:

1. the pretrained_checkpoint_config (if defined) or
2. the user uses existing model compression techniques to create a student

In the 2. case we might want a warmup block as shown in the above configs to allow the layers to realign its weights before going to the OPD process to prevent any rubbish token generation by the student during the forward loop

OPD

Once the student is initialized (by any of the paths above), we use a separate OPD flow. Because OPD requires the student to generate text, we introduce a specialized training loop wrapper. we use an OPDTrainer wrapper that performs the standard OPD flow from the paper (student generates a rollout, teacher scores it, top-K token selection, masked KL loss).

Vignettes

Here is how the user would interact with this design in their training script:

Using configs for everything

import deepspeed
from deepspeed.compression.opd import OPDTrainer

ds_config = "ds_config.json"

opd_trainer = OPDTrainer.from_config(ds_config)

here we can also allow the user to pass student and teacher directly and not to the config

opd_trainer.train()
Allow passing different student/teacher model to the trainer for flexibility

opd_trainer = OPDTrainer(
    student_model=student_model,
    teacher_model=teacher_model,
    ds_config="ds_config.json",
    train_dataset=dataset
)
opd_trainer.train()

Here, we in case we allow the current model compression techniques to be defined as student and we need a warmup loop - train should handle that after looking at the configs

Or do you prefer that we dont have a trainer and the loop is written by the user themselves?

@phoeenniixx Thanks for providing your thoughts on OPD Trainer design. We have merged the PR last month for OPD trainer, leveraging DeepSpeed's own HybridEngine rollout. IMO key bottleneck for OPD trainer is not config, but the rollout engine. We decided to go unified OPD framework so everything is contained within DeepSpeed. One question maybe you can help think about is how to improve OPD trainer performance. Especially the way to improve our HybridEngine performance.

[phoeenniixx]

Thanks for the reply @PKUWZP! I didnt find any OPD trainer implementation in compression module, I might have missed it. Can you please direct me to the implementation? I want to better understand how its implemented and how it uses the HybridEngine before making any comments :)

EDIT: I found this PR #8027, I understand I was looking at the wrong directory.. Thanks for all you help!

Just a question out of curiosity:
Does the engine allow any student model (like using a model compression method on the teacher to create a student) or is it completely same as the paper- we take pre-defined student-teacher pair?

[PKUWZP]

Thanks for the reply @PKUWZP! I didnt find any OPD trainer implementation in compression module, I might have missed it. Can you please direct me to the implementation? I want to better understand how its implemented and how it uses the HybridEngine before making any comments :)

EDIT: I found this PR #8027, I understand I was looking at the wrong directory.. Thanks for all you help!

Just a question out of curiosity: Does the engine allow any student model (like using a model compression method on the teacher to create a student) or is it completely same as the paper- we take pre-defined student-teacher pair?

The HybridEngine should be flexible supporting any student model. Assuming you are talking about the general model compression methods, OPD is actually one of the model compression techniques (the other techniques might be quantization/pruning etc.). The paper is just an example of it.

[phoeenniixx]

The HybridEngine should be flexible supporting any student model. Assuming you are talking about the general model compression methods, OPD is actually one of the model compression techniques (the other techniques might be quantization/pruning etc.). The paper is just an example of it.

I see, thanks! This means if I want to use Qwen3.6-27B as a teacher, but I perform layer reduction on this same model to create a student with 8B params (just an example), then I can use OPD on these models?
Do you think OPD will be effective out of the box in this scenario?

[PKUWZP]

The HybridEngine should be flexible supporting any student model. Assuming you are talking about the general model compression methods, OPD is actually one of the model compression techniques (the other techniques might be quantization/pruning etc.). The paper is just an example of it.

I see, thanks! This means if I want to use Qwen3.6-27B as a teacher, but I perform layer reduction on this same model to create a student with 8B params (just an example), then I can use OPD on these models? Do you think OPD will be effective out of the box in this scenario?

I don't think you need layer reduction from the teacher model to create student models. OPD by design should work out of the box, it has nothing to do with how the student models are created.