[WIP] Add length prediction task for multi-task learning

[Copilot]

Thanks for asking me to work on this. I will get started on it and keep this PR's description up to date as I form a plan and make progress.

Original prompt

目标

在现有的 EAGLE 训练代码中添加输出长度预测任务，实现多任务学习。

具体需求

1. 修改 eagle/traineagle3/cnets.py

1.1 在 Model 类的 __init__ 方法中添加长度预测头

# 在 self.lm_head 定义之后添加
self.length_head = nn.Sequential(
    nn.Linear(config.hidden_size, config.hidden_size // 2),
    nn.ReLU(),
    nn.Dropout(0.1),
    nn.Linear(config.hidden_size // 2, 1)  # 回归任务，预测长度
)

1.2 修改 dataprepare 方法，添加真实长度标签计算

在返回语句之前添加：

# 计算每个样本的有效输出长度
with torch.no_grad():
    # loss_mask shape: [batch_size, seq_len, 1]
    # 统计每个样本中 loss_mask=1 的位置数量
    real_output_lengths = loss_mask.squeeze(-1).sum(dim=1).float()  # [batch_size]

修改返回语句为：

return hidden_states, target, loss_mask, input_ids, real_output_lengths

1.3 修改 forward 方法

修改数据准备调用：

hidden_states, target, loss_mask, input_ids, real_lengths = self.dataprepare(input_ids, attention_mask, loss_mask)

在循环中添加长度预测：
在 for idx in range(self.length): 循环内，在计算 hidden_states_out 之后添加：

# 长度预测分支（只在最后一层预测）
if last:
    # 使用序列的平均表示进行长度预测
    pooled_hidden = hidden_states_out.mean(dim=1)  # [batch_size, hidden_size]
    length_pred = self.length_head(pooled_hidden).squeeze(-1)  # [batch_size]
    
    # 计算长度预测损失
    length_loss = nn.MSELoss()(length_pred, real_lengths)
else:
    length_loss = torch.tensor(0.0, device=hidden_states.device)

修改返回语句：
在循环结束后修改：

return plosses, vlosses, acces, length_loss

2. 修改 eagle/traineagle3/main.py

2.1 修改训练循环中的前向传播调用

找到这两处调用：

plosses, vlosses, acces = model_engine(...)

修改为：

plosses, vlosses, acces, length_loss = model_engine(...)

2.2 修改损失计算

在训练循环中，找到：

ploss_weight = [0.8 ** i for i in range(len(plosses))]
ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))])
loss = ploss

修改为：

ploss_weight = [0.8 ** i for i in range(len(plosses))]
ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))])

# 添加长度预测损失，权重可调
length_loss_weight = 0.1  # 可以调整这个权重
loss = ploss + length_loss_weight * length_loss

2.3 添加 wandb 日志记录

在训练循环的日志部分添加：

if global_rank == 0:
    logdict = {
        "train/lr": optimizer.optimizer.param_groups[0]["lr"],
        "train/length_loss": length_loss.item()
    }
    for i in range(len(plosses)):
        logdict[f"train/ploss_{i}"] = plosses[i].item()
    for i in range(len(acces)):
        logdict[f"train/acc_{i}"] = acces[i]
    wandb.log(logdict)

在测试循环中也添加相应的日志记录：

# 在测试循环外初始化
epoch_length_losses = []

# 在测试循环内
epoch_length_losses.append(length_loss.item())

# 在测试循环后
avg_length_loss = torch.tensor(epoch_length_losses).cuda().mean()
deepspeed.comm.all_reduce(avg_length_loss, op=deepspeed.comm.ReduceOp.AVG)
if global_rank == 0:
    wandb.log({"test/length_loss": avg_length_loss.item()})
    print(f"Test Epoch [{epoch + 1}/{num_epochs}], Length Loss: {avg_length_loss.item():.4f}")

3. 配置文件更新（可选）

如果需要，可以在 configs.py 或配置文件中添加：

length_loss_weight: float = 0.1
use_length_prediction: bool = True

实现要点

1. 长度标签来源：使用 loss_mask 统计每个样本的实际输出token数量
2. 预测时机：只在最后一次迭代时进行长度预测，减少计算开销
3. 特征提取：使用序列的平均池化表示作为长度预测的输入
4. 损失平衡：使用可调权重平衡token预测损失和长度预测损失
5. 梯度兼容：确保与现有的梯度检查点机制兼容

预期效果

• 模型同时学习预测下一个token和预测输出总长度
• 长度预测任务可能帮助模型更好地理解序列结构
• 可以在推理时使用长度预测来优化生成策略

This pull request was created as a result of the following prompt from Copilot chat.

目标

在现有的 EAGLE 训练代码中添加输出长度预测任务，实现多任务学习。

具体需求

1. 修改 eagle/traineagle3/cnets.py

1.1 在 Model 类的 __init__ 方法中添加长度预测头

# 在 self.lm_head 定义之后添加
self.length_head = nn.Sequential(
    nn.Linear(config.hidden_size, config.hidden_size // 2),
    nn.ReLU(),
    nn.Dropout(0.1),
    nn.Linear(config.hidden_size // 2, 1)  # 回归任务，预测长度
)

1.2 修改 dataprepare 方法，添加真实长度标签计算

在返回语句之前添加：

# 计算每个样本的有效输出长度
with torch.no_grad():
    # loss_mask shape: [batch_size, seq_len, 1]
    # 统计每个样本中 loss_mask=1 的位置数量
    real_output_lengths = loss_mask.squeeze(-1).sum(dim=1).float()  # [batch_size]

修改返回语句为：

return hidden_states, target, loss_mask, input_ids, real_output_lengths

1.3 修改 forward 方法

修改数据准备调用：

hidden_states, target, loss_mask, input_ids, real_lengths = self.dataprepare(input_ids, attention_mask, loss_mask)

在循环中添加长度预测：
在 for idx in range(self.length): 循环内，在计算 hidden_states_out 之后添加：

# 长度预测分支（只在最后一层预测）
if last:
    # 使用序列的平均表示进行长度预测
    pooled_hidden = hidden_states_out.mean(dim=1)  # [batch_size, hidden_size]
    length_pred = self.length_head(pooled_hidden).squeeze(-1)  # [batch_size]
    
    # 计算长度预测损失
    length_loss = nn.MSELoss()(length_pred, real_lengths)
else:
    length_loss = torch.tensor(0.0, device=hidden_states.device)

修改返回语句：
在循环结束后修改：

return plosses, vlosses, acces, length_loss

2. 修改 eagle/traineagle3/main.py

2.1 修改训练循环中的前向传播调用

找到这两处调用：

plosses, vlosses, acces = model_engine(...)

修改为：

plosses, vlosses, acces, length_loss = model_engine(...)

2.2 修改损失计算

在训练循环中，找到：

ploss_weight = [0.8 ** i for i in range(len(plosses))]
ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))])
loss = ploss

修改为：

ploss_weight = [0.8 ** i for i in range(len(plosses))]
ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))])

# 添加长度预测损失，权重可调
length_loss_weight = 0.1  # 可以调整这个权重
loss = ploss + length_loss_weight * length_loss

2.3 添加 wandb 日志记录

在训练循环的日志部分添加：

if global_rank == 0:
    logdict = {
        "train/lr": optimizer.optimizer.param_groups[0]["lr"],
        "train/length_loss": length_loss.item()
    }
    for i in range(len(plosses)):
        logdict[f"train/ploss_{i}"] = plosses[i].item()
    for i in range(len(acces)):
        logdict[f"train/acc_{i}"] = acces[i]
    wandb.log(logdict)

在测试循环中也添加相应的日志记录：

# 在测试循环外初始化
epoch_length_losses = []

# 在测试循环内
epoch_length_losses.append(length_loss.item())

# 在测试循环后
avg_length_loss = torch.tensor(epoch_length_losses).cuda().mean()
deepspeed.comm.all_reduce(avg_length_loss, op=deepspeed.comm.ReduceOp.AVG)
if global_rank == 0:
    wandb.log({"test/length_loss": avg_length_loss.item()})
    print(f"Test Epoch [{epoch + 1}/{num_epochs}], Length Loss: {avg_length_loss.item():.4f}")

3. 配置文件更新（可选）

如果需要，可以在 configs.py 或配置文件中添加：

length_loss_weight: float = 0.1
use_length_prediction: bool = True

实现要点

1. 长度标签来源：使用 loss_mask 统计每个样本的实际输出token数量
2. 预测时机：只在最后一次迭代时进行长度预测，减少计算开销
3. 特征提取：使用序列的平均池化表示作为长度预测的输入
4. 损失平衡：使用可调权重平衡token预测损失和长度预测损失
5. 梯度兼容：确保与现有的梯度检查点机制兼容

预期效果

• 模型同时学习预测下一个token和预测输出总长度
• 长度预测任务可能帮助模型更好地理解序列结构
• 可以在推理时使用长度预测来优化生成策略

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.