eval_step no backprop and new state compared with train_step.
(model, config, state, data, dropout_rng)
| 351 | |
| 352 | |
| 353 | def eval_step(model, config, state, data, dropout_rng): |
| 354 | """eval_step no backprop and new state compared with train_step.""" |
| 355 | |
| 356 | reference_params, extra_dpo_args, _loss_fn = [], [], loss_fn |
| 357 | if config.use_dpo: |
| 358 | state, reference_params = _split_dpo_state(state) |
| 359 | extra_dpo_args = [reference_params] |
| 360 | _loss_fn = dpo_loss_fn |
| 361 | |
| 362 | eval_loss_fn = functools.partial(_loss_fn, model, config, data, dropout_rng, is_train=False) |
| 363 | loss, aux = eval_loss_fn(state.params, *extra_dpo_args) |
| 364 | |
| 365 | mtp_acceptance_rate = 0.0 |
| 366 | if config.mtp_eval_target_module > 0: |
| 367 | mtp_acceptance_rate = calculate_mtp_acceptance_rate(aux["intermediate_outputs"], config) |
| 368 | |
| 369 | total_loss = aux["total_loss"] |
| 370 | total_weights = aux["total_weights"] |
| 371 | moe_lb_loss = aux["moe_lb_loss"] |
| 372 | mtp_loss = aux["mtp_loss"] |
| 373 | metrics = { |
| 374 | "scalar": { |
| 375 | "evaluation/loss": loss, |
| 376 | "evaluation/total_loss": total_loss, |
| 377 | "evaluation/total_weights": total_weights, |
| 378 | "evaluation/moe_lb_loss": moe_lb_loss, |
| 379 | "evaluation/mtp_loss": mtp_loss, |
| 380 | "evaluation/mtp_acceptance_rate_percent": mtp_acceptance_rate, |
| 381 | }, |
| 382 | } |
| 383 | if config.use_dpo: |
| 384 | metrics["scalar"]["evaluation/dpo_reward_accuracy"] = aux["reward_accuracy"] |
| 385 | |
| 386 | return metrics |
| 387 | |
| 388 | |
| 389 | def train_loop(config, recorder, state=None): |
nothing calls this directly
no test coverage detected