Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / OURS /src /rl /engine.py

EmergentTwilight

20 days ago

download

raw

13.2 kB

"""One fresh-group actor update for GRPO v0."""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch

from shared.training import assert_finite_training_state

from .distributed import DistributedContext
from .objective import multimodal_grpo_loss
from .policy import policy_replay_mismatch


@dataclass(frozen=True)
class PolicyRecord:
    text_current_log_probs: torch.Tensor
    text_old_log_probs: torch.Tensor
    text_reference_log_probs: torch.Tensor
    visual_current_factor_log_probs: torch.Tensor
    visual_old_factor_log_probs: torch.Tensor
    visual_reference_factor_log_probs: torch.Tensor
    text_replay_log_probs: torch.Tensor
    visual_replay_factor_log_probs: torch.Tensor
    advantage: float
    behavior_log_prob_max_abs_error: float
    text_behavior_log_prob_max_abs_error: float
    visual_behavior_log_prob_max_abs_error: float
    text_behavior_log_prob_mean_abs_error: float
    text_behavior_log_prob_p99_abs_error: float
    visual_behavior_log_prob_mean_abs_error: float
    visual_behavior_log_prob_p99_abs_error: float
    visual_behavior_log_prob_step_signed_mean_abs_p99: float
    text_ratio_mean: float
    visual_factor_ratio_mean: float
    text_clip_fraction: float
    visual_factor_clip_fraction: float
    reference_jsd: float
    autonomous_ar_steps: int


@dataclass(frozen=True)
class ActorUpdateResult:
    successful: bool
    zero_gradient: bool
    grad_norm: float
    trajectories: int
    statistics: dict[str, float]
    policy_records: tuple[PolicyRecord, ...]


def logical_batch_replay_statistics(
    records: list[PolicyRecord],
) -> dict[str, float]:
    """Combine stable trajectory guards with global action percentiles."""

    if not records:
        raise ValueError("replay guards require policy records")

    def population(replay_name: str, old_name: str, *, factorized: bool):
        pairs = [
            (getattr(record, replay_name), getattr(record, old_name))
            for record in records
            if getattr(record, replay_name).numel()
        ]
        if not pairs:
            empty = torch.empty((0, 0) if factorized else (0,))
            return policy_replay_mismatch(empty, empty, factorized=factorized)
        replay, old = zip(*pairs, strict=True)
        return policy_replay_mismatch(
            torch.cat(replay), torch.cat(old), factorized=factorized
        )

    text = population(
        "text_replay_log_probs", "text_old_log_probs", factorized=False
    )
    visual = population(
        "visual_replay_factor_log_probs",
        "visual_old_factor_log_probs",
        factorized=True,
    )
    return {
        "behavior_log_prob_max_abs_error": max(
            text["max_abs"], visual["max_abs"]
        ),
        "text_behavior_log_prob_mean_abs_error": max(
            record.text_behavior_log_prob_mean_abs_error for record in records
        ),
        "text_behavior_log_prob_p99_abs_error": text["p99_abs"],
        "visual_behavior_log_prob_mean_abs_error": max(
            record.visual_behavior_log_prob_mean_abs_error for record in records
        ),
        "visual_behavior_log_prob_p99_abs_error": visual["p99_abs"],
        "visual_behavior_log_prob_step_signed_mean_abs_p99": visual[
            "step_signed_mean_abs_p99"
        ],
    }


def replay_guard_violations(
    statistics: dict[str, float], config
) -> dict[str, tuple[float, float]]:
    limits = {
        "catastrophic_max": (
            "behavior_log_prob_max_abs_error",
            config.behavior.replay_catastrophic_max_abs_limit,
        ),
        "text_mean_abs": (
            "text_behavior_log_prob_mean_abs_error",
            config.behavior.replay_text_mean_abs_limit,
        ),
        "text_p99_abs": (
            "text_behavior_log_prob_p99_abs_error",
            config.behavior.replay_text_p99_abs_limit,
        ),
        "visual_mean_abs": (
            "visual_behavior_log_prob_mean_abs_error",
            config.behavior.replay_visual_mean_abs_limit,
        ),
        "visual_p99_abs": (
            "visual_behavior_log_prob_p99_abs_error",
            config.behavior.replay_visual_p99_abs_limit,
        ),
        "visual_step_signed_mean_abs_p99": (
            "visual_behavior_log_prob_step_signed_mean_abs_p99",
            config.behavior.replay_visual_step_signed_mean_abs_p99_limit,
        ),
    }
    return {
        name: (statistics[key], limit)
        for name, (key, limit) in limits.items()
        if statistics[key] > limit
    }


def execute_actor_update(
    *,
    groups,
    policy_scorer,
    config,
    parameters: list[torch.nn.Parameter],
    actor,
    optimizer,
    scheduler,
    parallel: DistributedContext | None = None,
) -> ActorUpdateResult:
    parallel = parallel or DistributedContext()
    trajectories = sum(len(group.trajectories) for group in groups)
    if not trajectories or not parameters:
        raise ValueError("actor update requires trajectories and parameters")
    if parallel.world_size > trajectories:
        raise ValueError("actor update requires at least one trajectory per worker")
    optimizer.zero_grad(set_to_none=True)
    finite = True
    statistics: dict[str, torch.Tensor] = {}
    indexed_records = []
    local_error = None
    flat_index = -1
    try:
        for group in groups:
            for index, scored in enumerate(group.trajectories):
                flat_index += 1
                if flat_index % parallel.world_size != parallel.rank:
                    continue
                policy = policy_scorer.score(
                    scored.trajectory,
                    advantage=group.advantages[index],
                )
                output = multimodal_grpo_loss(
                    [policy.inputs],
                    clip_low=config.objective.clip_lower,
                    clip_high=config.objective.clip_upper,
                    reference_beta=config.reference.beta,
                )
                finite &= bool(torch.isfinite(output.loss.detach()))
                (output.loss / trajectories).backward()
                for key, value in output.statistics.items():
                    statistics[key] = statistics.get(key, 0.0) + value.detach()
                inputs = policy.inputs
                indexed_records.append(
                    (
                        flat_index,
                        PolicyRecord(
                            text_current_log_probs=(
                                inputs.text_current_log_probs.detach().cpu()
                            ),
                            text_old_log_probs=inputs.text_old_log_probs.detach().cpu(),
                            text_reference_log_probs=(
                                policy.text_reference_log_probs.cpu()
                            ),
                            visual_current_factor_log_probs=(
                                inputs.visual_current_factor_log_probs.detach().cpu()
                            ),
                            visual_old_factor_log_probs=(
                                inputs.visual_old_factor_log_probs.detach().cpu()
                            ),
                            visual_reference_factor_log_probs=(
                                policy.visual_reference_factor_log_probs.cpu()
                            ),
                            text_replay_log_probs=policy.text_replay_log_probs.cpu(),
                            visual_replay_factor_log_probs=(
                                policy.visual_replay_factor_log_probs.cpu()
                            ),
                            advantage=float(inputs.advantage.detach().cpu()),
                            behavior_log_prob_max_abs_error=(
                                policy.behavior_log_prob_max_abs_error
                            ),
                            text_behavior_log_prob_max_abs_error=(
                                policy.text_behavior_log_prob_max_abs_error
                            ),
                            visual_behavior_log_prob_max_abs_error=(
                                policy.visual_behavior_log_prob_max_abs_error
                            ),
                            text_behavior_log_prob_mean_abs_error=(
                                policy.text_behavior_log_prob_mean_abs_error
                            ),
                            text_behavior_log_prob_p99_abs_error=(
                                policy.text_behavior_log_prob_p99_abs_error
                            ),
                            visual_behavior_log_prob_mean_abs_error=(
                                policy.visual_behavior_log_prob_mean_abs_error
                            ),
                            visual_behavior_log_prob_p99_abs_error=(
                                policy.visual_behavior_log_prob_p99_abs_error
                            ),
                            visual_behavior_log_prob_step_signed_mean_abs_p99=(
                                policy.visual_behavior_log_prob_step_signed_mean_abs_p99
                            ),
                            text_ratio_mean=float(
                                output.statistics["text_ratio_mean"].cpu()
                            ),
                            visual_factor_ratio_mean=float(
                                output.statistics["visual_factor_ratio_mean"].cpu()
                            ),
                            text_clip_fraction=float(
                                output.statistics["text_clip_fraction"].cpu()
                            ),
                            visual_factor_clip_fraction=float(
                                output.statistics["visual_factor_clip_fraction"].cpu()
                            ),
                            reference_jsd=float(
                                output.statistics["reference_jsd"].cpu()
                            ),
                            autonomous_ar_steps=int(
                                output.statistics["autonomous_ar_steps"].cpu()
                            ),
                        ),
                    )
                )
    except Exception as error:
        if not parallel.enabled:
            raise
        local_error = f"rank {parallel.rank}: {type(error).__name__}: {error}"

    payloads = parallel.all_gather_object(
        {
            "error": local_error,
            "finite": finite,
            "statistics": {
                key: float(value.float().cpu()) for key, value in statistics.items()
            },
            "records": indexed_records,
        }
    )
    errors = [payload["error"] for payload in payloads if payload["error"]]
    if errors:
        optimizer.zero_grad(set_to_none=True)
        raise RuntimeError("distributed actor update failed: " + "; ".join(errors))

    finite = all(payload["finite"] for payload in payloads)
    summed_statistics: dict[str, float] = {}
    for payload in payloads:
        for key, value in payload["statistics"].items():
            summed_statistics[key] = summed_statistics.get(key, 0.0) + value
    ordered_records = sorted(
        (record for payload in payloads for record in payload["records"]),
        key=lambda item: item[0],
    )
    if [index for index, _record in ordered_records] != list(range(trajectories)):
        raise RuntimeError("distributed policy records do not cover the global group")
    records = [record for _index, record in ordered_records]
    replay_statistics = logical_batch_replay_statistics(records)
    failed_replay_guards = replay_guard_violations(replay_statistics, config)
    if failed_replay_guards:
        optimizer.zero_grad(set_to_none=True)
        raise RuntimeError(
            "logical-batch sampling/full-prefix replay mismatch exceeded guard: "
            + ", ".join(
                f"{name}={value:.6g}>{limit:.6g}"
                for name, (value, limit) in failed_replay_guards.items()
            )
        )
    parallel.synchronize_gradients(parameters)

    gradients = [parameter.grad for parameter in parameters if parameter.grad is not None]
    grad_norm = torch.nn.utils.clip_grad_norm_(
        parameters, config.optimizer.max_grad_norm
    )
    grad_norm_value = float(grad_norm.detach().float().cpu())
    zero_gradient = not gradients or grad_norm_value == 0.0
    successful = bool(finite and math.isfinite(grad_norm_value))
    if successful:
        optimizer.step()
        assert_finite_training_state(parameters, optimizer)
        scheduler.step()
    optimizer.zero_grad(set_to_none=True)

    summarized = {}
    for key, value in summed_statistics.items():
        summarized[key] = (
            value if key.endswith("_steps") else value / trajectories
        )
    summarized.update(replay_statistics)
    return ActorUpdateResult(
        successful=successful,
        zero_gradient=zero_gradient,
        grad_norm=grad_norm_value,
        trajectories=trajectories,
        statistics=summarized,
        policy_records=tuple(records),
    )


__all__ = [
    "ActorUpdateResult",
    "execute_actor_update",
    "logical_batch_replay_statistics",
    "PolicyRecord",
    "replay_guard_violations",
]

Xet Storage Details

Size:: 13.2 kB
Xet hash:: dc0816da2d97f8ff564afa3d712218b7472f6ab90446d0605635b58bfcc262ab

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""One fresh-group actor update for GRPO v0."""

	from __future__ import annotations

	import math
	from dataclasses import dataclass

	import torch

	from shared.training import assert_finite_training_state

	from .distributed import DistributedContext
	from .objective import multimodal_grpo_loss
	from .policy import policy_replay_mismatch


	@dataclass(frozen=True)
	class PolicyRecord:
	text_current_log_probs: torch.Tensor
	text_old_log_probs: torch.Tensor
	text_reference_log_probs: torch.Tensor
	visual_current_factor_log_probs: torch.Tensor
	visual_old_factor_log_probs: torch.Tensor
	visual_reference_factor_log_probs: torch.Tensor
	text_replay_log_probs: torch.Tensor
	visual_replay_factor_log_probs: torch.Tensor
	advantage: float
	behavior_log_prob_max_abs_error: float
	text_behavior_log_prob_max_abs_error: float
	visual_behavior_log_prob_max_abs_error: float
	text_behavior_log_prob_mean_abs_error: float
	text_behavior_log_prob_p99_abs_error: float
	visual_behavior_log_prob_mean_abs_error: float
	visual_behavior_log_prob_p99_abs_error: float
	visual_behavior_log_prob_step_signed_mean_abs_p99: float
	text_ratio_mean: float
	visual_factor_ratio_mean: float
	text_clip_fraction: float
	visual_factor_clip_fraction: float
	reference_jsd: float
	autonomous_ar_steps: int


	@dataclass(frozen=True)
	class ActorUpdateResult:
	successful: bool
	zero_gradient: bool
	grad_norm: float
	trajectories: int
	statistics: dict[str, float]
	policy_records: tuple[PolicyRecord, ...]


	def logical_batch_replay_statistics(
	records: list[PolicyRecord],
	) -> dict[str, float]:
	"""Combine stable trajectory guards with global action percentiles."""

	if not records:
	raise ValueError("replay guards require policy records")

	def population(replay_name: str, old_name: str, *, factorized: bool):
	pairs = [
	(getattr(record, replay_name), getattr(record, old_name))
	for record in records
	if getattr(record, replay_name).numel()
	]
	if not pairs:
	empty = torch.empty((0, 0) if factorized else (0,))
	return policy_replay_mismatch(empty, empty, factorized=factorized)
	replay, old = zip(*pairs, strict=True)
	return policy_replay_mismatch(
	torch.cat(replay), torch.cat(old), factorized=factorized
	)

	text = population(
	"text_replay_log_probs", "text_old_log_probs", factorized=False
	)
	visual = population(
	"visual_replay_factor_log_probs",
	"visual_old_factor_log_probs",
	factorized=True,
	)
	return {
	"behavior_log_prob_max_abs_error": max(
	text["max_abs"], visual["max_abs"]
	),
	"text_behavior_log_prob_mean_abs_error": max(
	record.text_behavior_log_prob_mean_abs_error for record in records
	),
	"text_behavior_log_prob_p99_abs_error": text["p99_abs"],
	"visual_behavior_log_prob_mean_abs_error": max(
	record.visual_behavior_log_prob_mean_abs_error for record in records
	),
	"visual_behavior_log_prob_p99_abs_error": visual["p99_abs"],
	"visual_behavior_log_prob_step_signed_mean_abs_p99": visual[
	"step_signed_mean_abs_p99"
	],
	}


	def replay_guard_violations(
	statistics: dict[str, float], config
	) -> dict[str, tuple[float, float]]:
	limits = {
	"catastrophic_max": (
	"behavior_log_prob_max_abs_error",
	config.behavior.replay_catastrophic_max_abs_limit,
	),
	"text_mean_abs": (
	"text_behavior_log_prob_mean_abs_error",
	config.behavior.replay_text_mean_abs_limit,
	),
	"text_p99_abs": (
	"text_behavior_log_prob_p99_abs_error",
	config.behavior.replay_text_p99_abs_limit,
	),
	"visual_mean_abs": (
	"visual_behavior_log_prob_mean_abs_error",
	config.behavior.replay_visual_mean_abs_limit,
	),
	"visual_p99_abs": (
	"visual_behavior_log_prob_p99_abs_error",
	config.behavior.replay_visual_p99_abs_limit,
	),
	"visual_step_signed_mean_abs_p99": (
	"visual_behavior_log_prob_step_signed_mean_abs_p99",
	config.behavior.replay_visual_step_signed_mean_abs_p99_limit,
	),
	}
	return {
	name: (statistics[key], limit)
	for name, (key, limit) in limits.items()
	if statistics[key] > limit
	}


	def execute_actor_update(
	*,
	groups,
	policy_scorer,
	config,
	parameters: list[torch.nn.Parameter],
	actor,
	optimizer,
	scheduler,
	parallel: DistributedContext \| None = None,
	) -> ActorUpdateResult:
	parallel = parallel or DistributedContext()
	trajectories = sum(len(group.trajectories) for group in groups)
	if not trajectories or not parameters:
	raise ValueError("actor update requires trajectories and parameters")
	if parallel.world_size > trajectories:
	raise ValueError("actor update requires at least one trajectory per worker")
	optimizer.zero_grad(set_to_none=True)
	finite = True
	statistics: dict[str, torch.Tensor] = {}
	indexed_records = []
	local_error = None
	flat_index = -1
	try:
	for group in groups:
	for index, scored in enumerate(group.trajectories):
	flat_index += 1
	if flat_index % parallel.world_size != parallel.rank:
	continue
	policy = policy_scorer.score(
	scored.trajectory,
	advantage=group.advantages[index],
	)
	output = multimodal_grpo_loss(
	[policy.inputs],
	clip_low=config.objective.clip_lower,
	clip_high=config.objective.clip_upper,
	reference_beta=config.reference.beta,
	)
	finite &= bool(torch.isfinite(output.loss.detach()))
	(output.loss / trajectories).backward()
	for key, value in output.statistics.items():
	statistics[key] = statistics.get(key, 0.0) + value.detach()
	inputs = policy.inputs
	indexed_records.append(
	(
	flat_index,
	PolicyRecord(
	text_current_log_probs=(
	inputs.text_current_log_probs.detach().cpu()
	),
	text_old_log_probs=inputs.text_old_log_probs.detach().cpu(),
	text_reference_log_probs=(
	policy.text_reference_log_probs.cpu()
	),
	visual_current_factor_log_probs=(
	inputs.visual_current_factor_log_probs.detach().cpu()
	),
	visual_old_factor_log_probs=(
	inputs.visual_old_factor_log_probs.detach().cpu()
	),
	visual_reference_factor_log_probs=(
	policy.visual_reference_factor_log_probs.cpu()
	),
	text_replay_log_probs=policy.text_replay_log_probs.cpu(),
	visual_replay_factor_log_probs=(
	policy.visual_replay_factor_log_probs.cpu()
	),
	advantage=float(inputs.advantage.detach().cpu()),
	behavior_log_prob_max_abs_error=(
	policy.behavior_log_prob_max_abs_error
	),
	text_behavior_log_prob_max_abs_error=(
	policy.text_behavior_log_prob_max_abs_error
	),
	visual_behavior_log_prob_max_abs_error=(
	policy.visual_behavior_log_prob_max_abs_error
	),
	text_behavior_log_prob_mean_abs_error=(
	policy.text_behavior_log_prob_mean_abs_error
	),
	text_behavior_log_prob_p99_abs_error=(
	policy.text_behavior_log_prob_p99_abs_error
	),
	visual_behavior_log_prob_mean_abs_error=(
	policy.visual_behavior_log_prob_mean_abs_error
	),
	visual_behavior_log_prob_p99_abs_error=(
	policy.visual_behavior_log_prob_p99_abs_error
	),
	visual_behavior_log_prob_step_signed_mean_abs_p99=(
	policy.visual_behavior_log_prob_step_signed_mean_abs_p99
	),
	text_ratio_mean=float(
	output.statistics["text_ratio_mean"].cpu()
	),
	visual_factor_ratio_mean=float(
	output.statistics["visual_factor_ratio_mean"].cpu()
	),
	text_clip_fraction=float(
	output.statistics["text_clip_fraction"].cpu()
	),
	visual_factor_clip_fraction=float(
	output.statistics["visual_factor_clip_fraction"].cpu()
	),
	reference_jsd=float(
	output.statistics["reference_jsd"].cpu()
	),
	autonomous_ar_steps=int(
	output.statistics["autonomous_ar_steps"].cpu()
	),
	),
	)
	)
	except Exception as error:
	if not parallel.enabled:
	raise
	local_error = f"rank {parallel.rank}: {type(error).__name__}: {error}"

	payloads = parallel.all_gather_object(
	{
	"error": local_error,
	"finite": finite,
	"statistics": {
	key: float(value.float().cpu()) for key, value in statistics.items()
	},
	"records": indexed_records,
	}
	)
	errors = [payload["error"] for payload in payloads if payload["error"]]
	if errors:
	optimizer.zero_grad(set_to_none=True)
	raise RuntimeError("distributed actor update failed: " + "; ".join(errors))

	finite = all(payload["finite"] for payload in payloads)
	summed_statistics: dict[str, float] = {}
	for payload in payloads:
	for key, value in payload["statistics"].items():
	summed_statistics[key] = summed_statistics.get(key, 0.0) + value
	ordered_records = sorted(
	(record for payload in payloads for record in payload["records"]),
	key=lambda item: item[0],
	)
	if [index for index, _record in ordered_records] != list(range(trajectories)):
	raise RuntimeError("distributed policy records do not cover the global group")
	records = [record for _index, record in ordered_records]
	replay_statistics = logical_batch_replay_statistics(records)
	failed_replay_guards = replay_guard_violations(replay_statistics, config)
	if failed_replay_guards:
	optimizer.zero_grad(set_to_none=True)
	raise RuntimeError(
	"logical-batch sampling/full-prefix replay mismatch exceeded guard: "
	+ ", ".join(
	f"{name}={value:.6g}>{limit:.6g}"
	for name, (value, limit) in failed_replay_guards.items()
	)
	)
	parallel.synchronize_gradients(parameters)

	gradients = [parameter.grad for parameter in parameters if parameter.grad is not None]
	grad_norm = torch.nn.utils.clip_grad_norm_(
	parameters, config.optimizer.max_grad_norm
	)
	grad_norm_value = float(grad_norm.detach().float().cpu())
	zero_gradient = not gradients or grad_norm_value == 0.0
	successful = bool(finite and math.isfinite(grad_norm_value))
	if successful:
	optimizer.step()
	assert_finite_training_state(parameters, optimizer)
	scheduler.step()
	optimizer.zero_grad(set_to_none=True)

	summarized = {}
	for key, value in summed_statistics.items():
	summarized[key] = (
	value if key.endswith("_steps") else value / trajectories
	)
	summarized.update(replay_statistics)
	return ActorUpdateResult(
	successful=successful,
	zero_gradient=zero_gradient,
	grad_norm=grad_norm_value,
	trajectories=trajectories,
	statistics=summarized,
	policy_records=tuple(records),
	)


	__all__ = [
	"ActorUpdateResult",
	"execute_actor_update",
	"logical_batch_replay_statistics",
	"PolicyRecord",
	"replay_guard_violations",
	]